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(