From 9ef1af9799060e5a285e1efd7ef439bde373b77d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?F=C3=A9lix=20Malfait?= Date: Fri, 3 Jul 2026 10:40:03 +0200 Subject: [PATCH] fix(ai): make stream claims atomic via conditional UPDATEs with claim-or-queue send (#22481) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Rationale `activeStreamId` is the mutex that guarantees one live stream per thread — but claiming it is a plain read-then-update. The resolver checks it, then `streamAgentChat` enqueues the job **before** writing the claim. Two racing sends both pass the check, both start jobs, and each job's `resetStreamState` wipes the other's Redis chunk list — tokens from two answers interleave into the visible message. The same window exists for retry vs. send, the queue drain vs. send, and `stopAgentChatStream`, which cleared the claim **unguarded** (`{ id, userWorkspaceId }`) and could wipe a newer stream's claim entirely. ## Why this is the root cause, not a symptom patch Ownership must live in the `activeStreamId` column regardless of any locking mechanism — the queue-behind gate, the thread DTO, and stop all read it. So the correct primitive is a single-row compare-and-set on that column: `UPDATE … WHERE "activeStreamId" IS NULL` checked via affected rows, claim **before** enqueue, release on enqueue failure. Every mutation of the claim is now guarded on the observed value. Alternatives evaluated and rejected: - **BullMQ jobId dedup by threadId**: the driver appends a `-${v4()}` suffix to custom ids and dedups via a non-atomic `getJobs(['waiting'])` scan that ignores active jobs — two racing sends still run concurrently, and it does nothing for stop/retry races. - **`SELECT FOR UPDATE` / Redis SETNX / advisory locks**: all add a second mechanism (transaction plumbing or a second source of truth) to protect a single-row write that Postgres can already do atomically. Path-specific claim predicates fall out naturally: send/drain claim with `pendingQuestionMessageId IS NULL`, retry claims with `lastStreamError IS NOT NULL` (and restores the error if its enqueue fails) — closing the double-retry race for free. ## User impact Double-send (impatient double-click, two tabs, retry racing a queued drain) can currently garble the assistant's answer with interleaved tokens from two model runs and strand one stream's claim. All of these become deterministic: exactly one winner streams; the loser queues politely. ## Test plan - [x] New claim spec: conditional claim before enqueue, race-loser queues, halted-backlog send queues at the back and kicks the drain front-first, enqueue-failure releases the claim - [x] Retry spec updated: rollback restores the prior `lastStreamError`; guarded shapes asserted - [ ] CI green https://claude.ai/code/session_01Lyi6zTema2FMVVh8MD6c38 --- _Generated by [Claude Code](https://claude.ai/code/session_01Lyi6zTema2FMVVh8MD6c38)_ Review in cubic --- .../ai/ai-chat/jobs/stream-agent-chat.job.ts | 2 +- .../ai-chat/resolvers/agent-chat.resolver.ts | 22 +- ...agent-chat-streaming.service.claim.spec.ts | 138 ++++++ ...agent-chat-streaming.service.retry.spec.ts | 18 +- .../services/agent-chat-streaming.service.ts | 417 +++++++++++------- .../ai/ai-chat/services/agent-chat.service.ts | 13 + .../ai-chat/utils/get-cancel-channel.util.ts | 4 +- 7 files changed, 448 insertions(+), 166 deletions(-) create mode 100644 packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/__tests__/agent-chat-streaming.service.claim.spec.ts 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 5ed99249ac..ee6871d9c9 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 @@ -69,7 +69,7 @@ export class StreamAgentChatJob { await this.eventPublisherService.resetStreamState(data.threadId); const abortController = new AbortController(); - const cancelChannel = getCancelChannel(data.threadId); + const cancelChannel = getCancelChannel(data.threadId, data.streamId); await this.cancelSubscriberService.subscribe(cancelChannel, () => { abortController.abort(); 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 c9208e06cc..1310c267a0 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 @@ -221,6 +221,16 @@ export class AgentChatResolver { fileAttachments: fileAttachments ?? undefined, }); + if (result.queued) { + await this.eventPublisherService.publish({ + threadId, + workspaceId: workspace.id, + event: { type: 'queue-updated' }, + }); + + return { messageId: result.messageId, queued: true }; + } + return { messageId: result.messageId, queued: false, @@ -360,11 +370,14 @@ export class AgentChatResolver { const redis = this.redisClientService.getClient(); - await redis.publish(getCancelChannel(threadId), 'cancel'); + await redis.publish( + getCancelChannel(threadId, thread.activeStreamId), + 'cancel', + ); await this.threadRepository.update( workspaceId, - { id: threadId, userWorkspaceId }, + { id: threadId, userWorkspaceId, activeStreamId: thread.activeStreamId }, { activeStreamId: null }, ); @@ -446,7 +459,10 @@ export class AgentChatResolver { const redis = this.redisClientService.getClient(); - await redis.publish(getCancelChannel(threadId), 'cancel'); + await redis.publish( + getCancelChannel(threadId, thread.activeStreamId), + 'cancel', + ); } @Mutation(() => Boolean) 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 new file mode 100644 index 0000000000..a63668f935 --- /dev/null +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/__tests__/agent-chat-streaming.service.claim.spec.ts @@ -0,0 +1,138 @@ +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 claim & reap', () => { + const workspace = { id: 'workspace-id' } as WorkspaceEntity; + + const idleThread = { + id: 'thread-id', + title: 'Thread', + conversationSize: 0, + activeStreamId: null, + lastStreamError: null, + pendingQuestionMessageId: null, + }; + + const buildService = ({ + thread = idleThread, + claimAffected = 1, + queuedMessages = [] as unknown[], + } = {}) => { + const publishedEvents: Array<{ type: string }> = []; + const threadRepository = { + findOne: jest.fn().mockResolvedValue(thread), + findOneOrFail: jest.fn().mockResolvedValue(thread), + update: jest.fn().mockResolvedValue({ affected: claimAffected }), + }; + const messageQueueService = { add: jest.fn().mockResolvedValue(undefined) }; + const agentChatService = { + addMessage: jest + .fn() + .mockResolvedValue({ id: 'user-message-id', turnId: 'turn-id' }), + notifyThreadActivityUpdated: jest.fn().mockResolvedValue(undefined), + getMessagesForThread: jest.fn().mockResolvedValue([]), + getQueuedMessages: jest.fn().mockResolvedValue(queuedMessages), + hasQueuedMessages: jest + .fn() + .mockImplementation(() => Promise.resolve(queuedMessages.length > 0)), + queueMessage: jest.fn().mockResolvedValue({ id: 'queued-message-id' }), + promoteQueuedMessage: jest.fn().mockResolvedValue('turn-id'), + deleteQueuedMessage: jest.fn().mockResolvedValue(true), + }; + const eventPublisherService = { + publish: jest.fn().mockImplementation(({ event }) => { + publishedEvents.push(event); + + return Promise.resolve(); + }), + resetStreamState: 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, + ); + + return { + service, + threadRepository, + messageQueueService, + agentChatService, + eventPublisherService, + publishedEvents, + }; + }; + + const sendArguments = { + threadId: 'thread-id', + userWorkspaceId: 'user-workspace-id', + workspace, + text: 'hello', + browsingContext: null, + }; + + describe('streamAgentChat', () => { + it('claims the thread conditionally before enqueueing', async () => { + const { service, threadRepository } = buildService(); + + const result = await service.streamAgentChat(sendArguments); + + expect(result.queued).toBe(false); + expect(threadRepository.update).toHaveBeenCalledWith( + 'workspace-id', + expect.objectContaining({ id: 'thread-id' }), + expect.objectContaining({ lastStreamError: null }), + ); + }); + + it('queues the message when another stream wins the claim race', async () => { + const { service, agentChatService, messageQueueService } = buildService({ + claimAffected: 0, + }); + + const result = await service.streamAgentChat(sendArguments); + + expect(result.queued).toBe(true); + expect(agentChatService.queueMessage).toHaveBeenCalled(); + expect(messageQueueService.add).not.toHaveBeenCalled(); + }); + + it('queues behind a halted backlog and kicks the drain from the front', async () => { + const { service, agentChatService } = buildService({ + queuedMessages: [ + { + id: 'older-queued-id', + parts: [{ type: 'text', textContent: 'first in line' }], + }, + ], + }); + + const result = await service.streamAgentChat(sendArguments); + + expect(result.queued).toBe(true); + expect(agentChatService.queueMessage).toHaveBeenCalled(); + expect(agentChatService.promoteQueuedMessage).toHaveBeenCalledWith( + expect.objectContaining({ messageId: 'older-queued-id' }), + ); + }); + + it('releases the claim when enqueueing the job fails', async () => { + const { service, threadRepository, messageQueueService } = buildService(); + + messageQueueService.add.mockRejectedValue(new Error('redis down')); + + await expect(service.streamAgentChat(sendArguments)).rejects.toThrow( + 'redis down', + ); + + expect(threadRepository.update).toHaveBeenLastCalledWith( + 'workspace-id', + { id: 'thread-id', activeStreamId: expect.any(String) }, + { activeStreamId: null }, + ); + }); + }); +}); diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/__tests__/agent-chat-streaming.service.retry.spec.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/__tests__/agent-chat-streaming.service.retry.spec.ts index 794f454fb1..cb1fa2bf35 100644 --- a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/__tests__/agent-chat-streaming.service.retry.spec.ts +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/__tests__/agent-chat-streaming.service.retry.spec.ts @@ -38,7 +38,7 @@ describe('AgentChatStreamingService.retryLastFailedTurn', () => { } = {}) => { const threadRepository = { findOne: jest.fn().mockResolvedValue(thread), - update: jest.fn().mockResolvedValue(undefined), + update: jest.fn().mockResolvedValue({ affected: 1 }), }; const messageQueueService = { add: jest.fn().mockResolvedValue(undefined) }; const agentChatService = { @@ -91,13 +91,13 @@ describe('AgentChatStreamingService.retryLastFailedTurn', () => { expect(messageQueueService.add).not.toHaveBeenCalled(); }); - it('rejects without clearing state when a newer message exists', async () => { + it('rejects and restores the error state when a newer message exists', async () => { const newerAssistantMessage = { ...userMessageEntity, id: 'newer-message-id', role: AgentMessageRole.ASSISTANT, } as unknown as AgentMessageEntity; - const { service, threadRepository } = buildService({ + const { service, threadRepository, messageQueueService } = buildService({ threadMessages: [userMessageEntity, newerAssistantMessage], }); @@ -106,7 +106,15 @@ describe('AgentChatStreamingService.retryLastFailedTurn', () => { ).rejects.toMatchObject({ code: AiExceptionCode.NO_FAILED_TURN_TO_RETRY, }); - expect(threadRepository.update).not.toHaveBeenCalled(); + expect(messageQueueService.add).not.toHaveBeenCalled(); + expect(threadRepository.update).toHaveBeenLastCalledWith( + 'workspace-id', + { id: 'thread-id', activeStreamId: expect.any(String) }, + { + activeStreamId: null, + lastStreamError: failedThread.lastStreamError, + }, + ); }); it('drops the failed output, re-enqueues the turn, and clears the error', async () => { @@ -134,7 +142,7 @@ describe('AgentChatStreamingService.retryLastFailedTurn', () => { ); expect(threadRepository.update).toHaveBeenCalledWith( 'workspace-id', - { id: 'thread-id' }, + expect.objectContaining({ id: 'thread-id' }), { activeStreamId: result.streamId, lastStreamError: null }, ); expect(result.messageId).toBe('user-message-id'); 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 b46ad4f97e..19ca567a61 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 @@ -8,7 +8,7 @@ import { } from 'twenty-shared/ai'; import { FileFolder } from 'twenty-shared/types'; import { isDefined } from 'twenty-shared/utils'; -import { In, Like } from 'typeorm'; +import { type FindOptionsWhere, In, IsNull, Like, Not } from 'typeorm'; import { FileEntity } from 'src/engine/core-modules/file/entities/file.entity'; import { FileUrlService } from 'src/engine/core-modules/file/file-url/file-url.service'; @@ -23,6 +23,7 @@ import { import { mapDBPartsToUIMessageParts } from 'src/engine/metadata-modules/ai/ai-agent-execution/utils/mapDBPartsToUIMessageParts'; import { type BrowsingContextType } from 'src/engine/metadata-modules/ai/ai-agent/types/browsingContext.type'; import { AgentChatThreadEntity } from 'src/engine/metadata-modules/ai/ai-chat/entities/agent-chat-thread.entity'; +import { type AgentChatThreadLastStreamError } from 'src/engine/metadata-modules/ai/ai-chat/types/agent-chat-thread-last-stream-error.type'; import { STREAM_AGENT_CHAT_JOB_NAME } from 'src/engine/metadata-modules/ai/ai-chat/jobs/stream-agent-chat-job-name.constant'; import { type StreamAgentChatJobData } from 'src/engine/metadata-modules/ai/ai-chat/jobs/stream-agent-chat-job.types'; import { AgentChatEventPublisherService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat-event-publisher.service'; @@ -61,6 +62,26 @@ export class AgentChatStreamingService { private readonly fileUrlService: FileUrlService, ) {} + private async tryClaimStream({ + threadId, + workspaceId, + streamId, + where, + }: { + threadId: string; + workspaceId: string; + streamId: string; + where: FindOptionsWhere; + }): Promise { + const claim = await this.threadRepository.update( + workspaceId, + { id: threadId, activeStreamId: IsNull(), ...where }, + { activeStreamId: streamId, lastStreamError: null }, + ); + + return claim.affected === 1; + } + async streamAgentChat({ threadId, userWorkspaceId, @@ -70,7 +91,10 @@ export class AgentChatStreamingService { modelId, messageId, fileAttachments, - }: StreamAgentChatOptions): Promise<{ streamId: string; messageId: string }> { + }: StreamAgentChatOptions): Promise< + | { queued: false; streamId: string; messageId: string } + | { queued: true; messageId: string } + > { const thread = await this.threadRepository.findOne(workspace.id, { where: { id: threadId, @@ -85,65 +109,100 @@ export class AgentChatStreamingService { ); } - const fileParts = await this.buildFilePartsFromAttachments( - fileAttachments, - workspace.id, - ); - - const userMessageParts: ExtendedUIMessagePart[] = [ - { type: 'text' as const, text }, - ...fileParts, - ]; - - const savedUserMessage = await this.agentChatService.addMessage({ + const hasQueuedBacklog = await this.agentChatService.hasQueuedMessages({ threadId, - id: messageId, - uiMessage: { - role: AgentMessageRole.USER, - parts: userMessageParts, - }, workspaceId: workspace.id, }); - await this.agentChatService.notifyThreadActivityUpdated({ - threadId, - userWorkspaceId, - workspaceId: workspace.id, - }); - - const previousMessages = await this.loadMessagesFromDB( - threadId, - userWorkspaceId, - workspace.id, - ); - const streamId = generateId(); - await this.messageQueueService.add( - STREAM_AGENT_CHAT_JOB_NAME, - { - threadId: thread.id, + const claimed = + !hasQueuedBacklog && + (await this.tryClaimStream({ + threadId, + workspaceId: workspace.id, streamId, + where: { pendingQuestionMessageId: IsNull() }, + })); + + if (!claimed) { + const queuedMessage = await this.agentChatService.queueMessage({ + threadId, + text, + id: messageId, + fileAttachments, + workspaceId: workspace.id, + userWorkspaceId, + }); + + if (hasQueuedBacklog) { + await this.flushNextQueuedMessage( + threadId, + userWorkspaceId, + workspace.id, + !!thread.title, + ); + } + + return { queued: true, messageId: queuedMessage.id }; + } + + try { + const fileParts = await this.buildFilePartsFromAttachments( + fileAttachments, + workspace.id, + ); + + const userMessageParts: ExtendedUIMessagePart[] = [ + { type: 'text' as const, text }, + ...fileParts, + ]; + + const savedUserMessage = await this.agentChatService.addMessage({ + threadId, + id: messageId, + uiMessage: { + role: AgentMessageRole.USER, + parts: userMessageParts, + }, + workspaceId: workspace.id, + }); + + await this.agentChatService.notifyThreadActivityUpdated({ + threadId, userWorkspaceId, workspaceId: workspace.id, - messages: previousMessages, - browsingContext, - modelId, - lastUserMessageText: text, - lastUserMessageParts: userMessageParts, - hasTitle: !!thread.title, - conversationSizeTokens: thread.conversationSize, - existingTurnId: savedUserMessage.turnId ?? undefined, - }, - ); + }); - await this.threadRepository.update( - workspace.id, - { id: thread.id }, - { activeStreamId: streamId, lastStreamError: null }, - ); + const previousMessages = await this.loadMessagesFromDB( + threadId, + userWorkspaceId, + workspace.id, + ); - return { streamId, messageId: savedUserMessage.id }; + await this.messageQueueService.add( + STREAM_AGENT_CHAT_JOB_NAME, + { + threadId: thread.id, + streamId, + userWorkspaceId, + workspaceId: workspace.id, + messages: previousMessages, + browsingContext, + modelId, + lastUserMessageText: text, + lastUserMessageParts: userMessageParts, + hasTitle: !!thread.title, + conversationSizeTokens: thread.conversationSize, + existingTurnId: savedUserMessage.turnId ?? undefined, + }, + ); + + return { queued: false, streamId, messageId: savedUserMessage.id }; + } catch (error) { + await this.releaseStreamClaim(threadId, workspace.id, streamId); + throw error; + } } async retryLastFailedTurn({ @@ -178,68 +237,85 @@ export class AgentChatStreamingService { ); } - const lastUserMessage = - await this.agentChatService.findLatestSentUserMessage({ - threadId, + const streamId = generateId(); + + const claimed = await this.tryClaimStream({ + threadId, + workspaceId: workspace.id, + streamId, + where: { lastStreamError: Not(IsNull()) }, + }); + + if (!claimed) { + throw new AiException( + 'There is no failed turn to retry on this thread', + AiExceptionCode.NO_FAILED_TURN_TO_RETRY, + ); + } + + try { + const lastUserMessage = + await this.agentChatService.findLatestSentUserMessage({ + threadId, + workspaceId: workspace.id, + }); + + if (!isDefined(lastUserMessage) || !isDefined(lastUserMessage.turnId)) { + throw new AiException( + 'There is no failed turn to retry on this thread', + AiExceptionCode.NO_FAILED_TURN_TO_RETRY, + ); + } + + await this.agentChatService.deleteAssistantMessagesForTurn({ + turnId: lastUserMessage.turnId, workspaceId: workspace.id, }); - if (!isDefined(lastUserMessage) || !isDefined(lastUserMessage.turnId)) { - throw new AiException( - 'There is no failed turn to retry on this thread', - AiExceptionCode.NO_FAILED_TURN_TO_RETRY, - ); - } - - await this.agentChatService.deleteAssistantMessagesForTurn({ - turnId: lastUserMessage.turnId, - workspaceId: workspace.id, - }); - - const messages = await this.loadMessagesFromDB( - threadId, - userWorkspaceId, - workspace.id, - ); - - const retriedMessage = messages[messages.length - 1]; - - if (!retriedMessage || retriedMessage.id !== lastUserMessage.id) { - throw new AiException( - 'There is no failed turn to retry on this thread', - AiExceptionCode.NO_FAILED_TURN_TO_RETRY, - ); - } - - const textPart = retriedMessage.parts.find((part) => part.type === 'text'); - - const streamId = generateId(); - - await this.messageQueueService.add( - STREAM_AGENT_CHAT_JOB_NAME, - { + const messages = await this.loadMessagesFromDB( threadId, - streamId, userWorkspaceId, - workspaceId: workspace.id, - messages, - browsingContext: null, - modelId, - lastUserMessageText: textPart?.text ?? '', - lastUserMessageParts: retriedMessage.parts, - hasTitle: !!thread.title, - conversationSizeTokens: thread.conversationSize, - existingTurnId: lastUserMessage.turnId, - }, - ); + workspace.id, + ); - await this.threadRepository.update( - workspace.id, - { id: threadId }, - { activeStreamId: streamId, lastStreamError: null }, - ); + const retriedMessage = messages[messages.length - 1]; - return { streamId, messageId: lastUserMessage.id }; + if (!retriedMessage || retriedMessage.id !== lastUserMessage.id) { + throw new AiException( + 'There is no failed turn to retry on this thread', + AiExceptionCode.NO_FAILED_TURN_TO_RETRY, + ); + } + + const textPart = retriedMessage.parts.find( + (part) => part.type === 'text', + ); + + await this.messageQueueService.add( + STREAM_AGENT_CHAT_JOB_NAME, + { + threadId, + streamId, + userWorkspaceId, + workspaceId: workspace.id, + messages, + browsingContext: null, + modelId, + lastUserMessageText: textPart?.text ?? '', + lastUserMessageParts: retriedMessage.parts, + hasTitle: !!thread.title, + conversationSizeTokens: thread.conversationSize, + existingTurnId: lastUserMessage.turnId, + }, + ); + + return { streamId, messageId: lastUserMessage.id }; + } catch (error) { + await this.releaseStreamClaim(threadId, workspace.id, streamId, { + lastStreamError: thread.lastStreamError, + }); + throw error; + } } async enqueueResumeStream({ @@ -340,66 +416,97 @@ export class AgentChatStreamingService { return; } - const turnId = await this.agentChatService.promoteQueuedMessage({ - messageId: nextQueued.id, + const streamId = generateId(); + + const claimed = await this.tryClaimStream({ threadId, workspaceId, + streamId, + where: { pendingQuestionMessageId: IsNull() }, }); - if (turnId === null) { + if (!claimed) { return; } - await this.eventPublisherService.publish({ - threadId, - workspaceId, - event: { type: 'queue-updated' }, - }); - - await this.eventPublisherService.publish({ - threadId, - workspaceId, - event: { type: 'message-persisted', messageId: nextQueued.id }, - }); - - const [uiMessages, thread] = await Promise.all([ - this.loadMessagesFromDB(threadId, userWorkspaceId, workspaceId), - this.threadRepository.findOneOrFail(workspaceId, { - where: { id: threadId }, - }), - ]); - - const streamId = generateId(); - - const lastUserMessageParts: ExtendedUIMessagePart[] = [ - ...(messageText !== '' - ? [{ type: 'text' as const, text: messageText }] - : []), - ...fileParts, - ]; - - await this.messageQueueService.add( - STREAM_AGENT_CHAT_JOB_NAME, - { + try { + const turnId = await this.agentChatService.promoteQueuedMessage({ + messageId: nextQueued.id, threadId, - streamId, - userWorkspaceId, workspaceId, - messages: uiMessages, - browsingContext: null, - lastUserMessageText: messageText, - lastUserMessageParts, - hasTitle, - conversationSizeTokens: thread.conversationSize, - existingTurnId: turnId, - }, - ); + }); - await this.threadRepository.update( - workspaceId, - { id: threadId }, - { activeStreamId: streamId, lastStreamError: null }, - ); + if (turnId === null) { + await this.releaseStreamClaim(threadId, workspaceId, streamId); + + return; + } + + await this.eventPublisherService.publish({ + threadId, + workspaceId, + event: { type: 'queue-updated' }, + }); + + await this.eventPublisherService.publish({ + threadId, + workspaceId, + event: { type: 'message-persisted', messageId: nextQueued.id }, + }); + + const [uiMessages, thread] = await Promise.all([ + this.loadMessagesFromDB(threadId, userWorkspaceId, workspaceId), + this.threadRepository.findOneOrFail(workspaceId, { + where: { id: threadId }, + }), + ]); + + const lastUserMessageParts: ExtendedUIMessagePart[] = [ + ...(messageText !== '' + ? [{ type: 'text' as const, text: messageText }] + : []), + ...fileParts, + ]; + + await this.messageQueueService.add( + STREAM_AGENT_CHAT_JOB_NAME, + { + threadId, + streamId, + userWorkspaceId, + workspaceId, + messages: uiMessages, + browsingContext: null, + lastUserMessageText: messageText, + lastUserMessageParts, + hasTitle, + conversationSizeTokens: thread.conversationSize, + existingTurnId: turnId, + }, + ); + } catch (error) { + await this.releaseStreamClaim(threadId, workspaceId, streamId); + throw error; + } + } + + private async releaseStreamClaim( + threadId: string, + workspaceId: string, + streamId: string, + restore?: { lastStreamError: AgentChatThreadLastStreamError | null }, + ): Promise { + await this.threadRepository + .update( + workspaceId, + { id: threadId, activeStreamId: streamId }, + { activeStreamId: null, ...restore }, + ) + .catch((error) => { + this.logger.error( + `Failed to release stream claim for thread ${threadId}: ${error instanceof Error ? error.message : String(error)}`, + ); + }); } private async loadMessagesFromDB( 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 7ab108692b..5a7ef7b3e4 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 @@ -408,6 +408,19 @@ export class AgentChatService { } as AgentMessageEntity; } + async hasQueuedMessages({ + threadId, + workspaceId, + }: { + threadId: string; + workspaceId: string; + }): Promise { + return this.messageRepository.existsBy(workspaceId, { + threadId, + status: AgentMessageStatus.QUEUED, + }); + } + async getQueuedMessages({ threadId, workspaceId, diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/utils/get-cancel-channel.util.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/utils/get-cancel-channel.util.ts index e161bc491c..ceff29b2c3 100644 --- a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/utils/get-cancel-channel.util.ts +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/utils/get-cancel-channel.util.ts @@ -1,2 +1,2 @@ -export const getCancelChannel = (threadId: string) => - `ai-stream:cancel:${threadId}`; +export const getCancelChannel = (threadId: string, streamId: string) => + `ai-stream:cancel:${threadId}:${streamId}`;