From 6ac8ebcd180707feb5f0e5f269b55bb1e732e351 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?F=C3=A9lix=20Malfait?= Date: Thu, 2 Jul 2026 21:01:32 +0200 Subject: [PATCH] test(ai): pin StreamAgentChatJob stream lifecycle with a reusable spec harness (#22479) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Rationale `StreamAgentChatJob` is the most failure-sensitive path in the AI chat stack — it coordinates the model stream, Redis event publishing, message persistence, the thread's stream claim, and the queued-message flush — and it had **zero unit coverage**. Both historical hangs lived here: - a throw before the model stream merges bypassed `onFinish` entirely and hung the job until the **10-minute BullMQ lock** expired (thread stuck the whole time); - a throw inside `onFinish` (persistence failure) left trailing chunks published with **no terminal event** — the client spinner ran forever. Production still shows this class is live: `Query read timeout` thrown from inside `handleStreamFinish` ([TWENTY-SERVER-GV7](https://twenty-v7.sentry.io/issues/TWENTY-SERVER-GV7)). ## Why this shape, not something else Tests-only PR, zero production risk. The fake chat stream mirrors the one AI SDK contract the job's coordination depends on — verified against the installed `ai@6.0.97` dist: `toUIMessageStream` converts mid-stream errors into error parts and **always** fires `onFinish` when the stream ends (`handleUIMessageStreamFinish` invokes it from both `flush()` and `cancel()`). Pinning that contract in the fake means a future SDK upgrade that breaks it fails these tests instead of production. Six tests pin current behavior: chunk ordering with `message-persisted` last, opaque error-chunk suppression, mid-stream failure persisting `lastStreamError` + releasing the claim, the two hang regressions above, and cancel skipping the queued flush. Three sibling PRs extend this exact spec file (missing-workspace routing, halted queue, and — later — auto-retry), which is why the harness lands first. ## User impact None directly; it makes the two worst historical user-facing hangs (10-minute dead thread, infinite spinner) regression-proof before the stuck-state fix series touches this code. ## Test plan - [x] 6 unit tests, no production code changed - [ ] CI green https://claude.ai/code/session_01Lyi6zTema2FMVVh8MD6c38 --- _Generated by [Claude Code](https://claude.ai/code/session_01Lyi6zTema2FMVVh8MD6c38)_ Review in cubic --- .../__tests__/stream-agent-chat.job.spec.ts | 318 ++++++++++++++++++ 1 file changed, 318 insertions(+) create mode 100644 packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/jobs/__tests__/stream-agent-chat.job.spec.ts diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/jobs/__tests__/stream-agent-chat.job.spec.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/jobs/__tests__/stream-agent-chat.job.spec.ts new file mode 100644 index 0000000000..44b38bdfb7 --- /dev/null +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/jobs/__tests__/stream-agent-chat.job.spec.ts @@ -0,0 +1,318 @@ +import { type UIMessageChunk } from 'ai'; + +import { type WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.entity'; +import { StreamAgentChatJob } from 'src/engine/metadata-modules/ai/ai-chat/jobs/stream-agent-chat.job'; +import { type StreamAgentChatJobData } from 'src/engine/metadata-modules/ai/ai-chat/jobs/stream-agent-chat-job.types'; + +type PublishedEvent = { type: string } & Record; + +const TEXT_CHUNKS: UIMessageChunk[] = [ + { type: 'start', messageId: 'assistant-message-id' }, + { type: 'start-step' }, + { type: 'text-start', id: 'text-1' }, + { type: 'text-delta', id: 'text-1', delta: 'Hello' }, + { type: 'text-end', id: 'text-1' }, +]; + +const RESPONSE_MESSAGE = { + role: 'assistant' as const, + parts: [{ type: 'text' as const, text: 'Hello' }], +}; + +const createFakeChatStream = ({ + chunks = TEXT_CHUNKS, + responseMessage = RESPONSE_MESSAGE, + midStreamError, + onFirstChunk, +}: { + chunks?: UIMessageChunk[]; + responseMessage?: typeof RESPONSE_MESSAGE; + midStreamError?: Error; + onFirstChunk?: () => void; +} = {}) => ({ + toUIMessageStream: (options: { + onError?: (error: unknown) => string; + onFinish?: (event: { + responseMessage: typeof RESPONSE_MESSAGE; + isAborted: boolean; + }) => Promise | void; + }) => + new ReadableStream({ + async start(controller) { + let isFirstChunk = true; + + for (const chunk of chunks) { + controller.enqueue(chunk); + + if (isFirstChunk) { + isFirstChunk = false; + onFirstChunk?.(); + } + } + + if (midStreamError) { + const errorText = options.onError?.(midStreamError) ?? ''; + + controller.enqueue({ type: 'error', errorText }); + } + + await options.onFinish?.({ responseMessage, isAborted: false }); + + controller.close(); + }, + }), +}); + +describe('StreamAgentChatJob', () => { + const workspace = { id: 'workspace-id' } as WorkspaceEntity; + + const jobData: StreamAgentChatJobData = { + threadId: 'thread-id', + streamId: 'stream-id', + userWorkspaceId: 'user-workspace-id', + workspaceId: 'workspace-id', + messages: [], + browsingContext: null, + lastUserMessageText: 'hello', + lastUserMessageParts: [{ type: 'text', text: 'hello' }], + hasTitle: true, + conversationSizeTokens: 0, + existingTurnId: 'turn-id', + }; + + const buildJob = ({ + workspaceFound = true, + chatStream = createFakeChatStream(), + streamChatRejection, + addMessageRejection, + }: { + workspaceFound?: boolean; + chatStream?: ReturnType; + streamChatRejection?: Error; + addMessageRejection?: Error; + } = {}) => { + const publishedEvents: PublishedEvent[] = []; + + const threadRepository = { + findOne: jest + .fn() + .mockResolvedValue({ id: 'thread-id', deletedAt: null }), + update: jest.fn().mockResolvedValue({ affected: 1 }), + }; + const workspaceRepository = { + findOne: jest.fn().mockResolvedValue(workspaceFound ? workspace : null), + }; + const agentChatService = { + addMessage: addMessageRejection + ? jest.fn().mockRejectedValue(addMessageRejection) + : jest.fn().mockResolvedValue({ id: 'assistant-message-id' }), + hasMessageById: jest.fn().mockResolvedValue(false), + generateTitleIfNeeded: jest.fn().mockResolvedValue(null), + notifyThreadUsageUpdated: jest.fn().mockResolvedValue(undefined), + }; + const chatExecutionService = { + streamChat: streamChatRejection + ? jest.fn().mockRejectedValue(streamChatRejection) + : jest.fn().mockResolvedValue({ + stream: chatStream, + modelConfig: { contextWindowTokens: 100000 }, + hasNoMoreAvailableCredits: () => false, + }), + }; + const eventPublisherService = { + resetStreamState: jest.fn().mockResolvedValue(undefined), + publish: jest + .fn() + .mockImplementation(({ event }: { event: PublishedEvent }) => { + publishedEvents.push(event); + + return Promise.resolve(); + }), + }; + const cancelCallbacks: Array<() => void> = []; + const cancelSubscriberService = { + subscribe: jest + .fn() + .mockImplementation((_channel: string, callback: () => void) => { + cancelCallbacks.push(callback); + + return Promise.resolve(); + }), + unsubscribe: jest.fn().mockResolvedValue(undefined), + }; + const agentChatStreamingService = { + flushNextQueuedMessage: jest.fn().mockResolvedValue(undefined), + }; + const job = new StreamAgentChatJob( + threadRepository as never, + workspaceRepository as never, + agentChatService as never, + chatExecutionService as never, + eventPublisherService as never, + cancelSubscriberService as never, + agentChatStreamingService as never, + ); + + return { + job, + publishedEvents, + threadRepository, + agentChatService, + eventPublisherService, + agentChatStreamingService, + cancelCallbacks, + }; + }; + + it('publishes all chunks in order with message-persisted last on success', async () => { + const { + job, + publishedEvents, + threadRepository, + agentChatService, + agentChatStreamingService, + } = buildJob(); + + await job.handle(jobData); + + const chunkEvents = publishedEvents.filter( + (event) => event.type === 'stream-chunk', + ); + + expect(chunkEvents).toHaveLength(TEXT_CHUNKS.length); + expect(publishedEvents[publishedEvents.length - 1]).toMatchObject({ + type: 'message-persisted', + }); + expect(agentChatService.addMessage).toHaveBeenCalledWith( + expect.objectContaining({ turnId: 'turn-id' }), + ); + expect(threadRepository.update).toHaveBeenCalledWith( + 'workspace-id', + { id: 'thread-id' }, + expect.objectContaining({ lastStreamError: null }), + ); + expect(threadRepository.update).toHaveBeenCalledWith( + 'workspace-id', + { id: 'thread-id', activeStreamId: 'stream-id' }, + { activeStreamId: null }, + ); + expect(agentChatStreamingService.flushNextQueuedMessage).toHaveBeenCalled(); + }); + + it('never publishes the opaque error chunk to subscribers', async () => { + const { job, publishedEvents } = buildJob({ + chatStream: createFakeChatStream({ + midStreamError: new Error('provider exploded'), + }), + }); + + await expect(job.handle(jobData)).rejects.toThrow('provider exploded'); + + const chunkTypes = publishedEvents + .filter((event) => event.type === 'stream-chunk') + .map((event) => (event.chunk as { type: string }).type); + + expect(chunkTypes).not.toContain('error'); + }); + + it('rejects, persists the error, and unblocks the thread when the model stream fails mid-stream', async () => { + const { job, publishedEvents, threadRepository } = buildJob({ + chatStream: createFakeChatStream({ + midStreamError: new Error('provider exploded'), + }), + }); + + await expect(job.handle(jobData)).rejects.toThrow('provider exploded'); + + const chunkEvents = publishedEvents.filter( + (event) => event.type === 'stream-chunk', + ); + + expect(chunkEvents).toHaveLength(TEXT_CHUNKS.length); + expect(publishedEvents[publishedEvents.length - 1]).toMatchObject({ + type: 'stream-error', + code: 'STREAM_EXECUTION_FAILED', + message: 'provider exploded', + }); + expect(publishedEvents.map((event) => event.type)).not.toContain( + 'message-persisted', + ); + expect(threadRepository.update).toHaveBeenCalledWith( + 'workspace-id', + { id: 'thread-id' }, + { + lastStreamError: expect.objectContaining({ + code: 'STREAM_EXECUTION_FAILED', + message: 'provider exploded', + }), + }, + ); + expect(threadRepository.update).toHaveBeenCalledWith( + 'workspace-id', + { id: 'thread-id', activeStreamId: 'stream-id' }, + { activeStreamId: null }, + ); + }); + + it('rejects promptly when execution setup throws instead of hanging until the queue lock expires', async () => { + const { job, publishedEvents, threadRepository } = buildJob({ + streamChatRejection: new Error('model resolution failed'), + }); + + await expect(job.handle(jobData)).rejects.toThrow( + 'model resolution failed', + ); + + expect(publishedEvents[publishedEvents.length - 1]).toMatchObject({ + type: 'stream-error', + message: 'model resolution failed', + }); + expect(threadRepository.update).toHaveBeenCalledWith( + 'workspace-id', + { id: 'thread-id', activeStreamId: 'stream-id' }, + { activeStreamId: null }, + ); + }); + + it('terminates the stream with an error when assistant persistence fails after draining chunks', async () => { + const { job, publishedEvents } = buildJob({ + addMessageRejection: new Error('insert failed'), + }); + + await expect(job.handle(jobData)).rejects.toThrow('insert failed'); + + const chunkEvents = publishedEvents.filter( + (event) => event.type === 'stream-chunk', + ); + + expect(chunkEvents).toHaveLength(TEXT_CHUNKS.length); + expect(publishedEvents[publishedEvents.length - 1]).toMatchObject({ + type: 'stream-error', + }); + expect(publishedEvents.map((event) => event.type)).not.toContain( + 'message-persisted', + ); + }); + + it('resolves without flushing the queue when the stream is cancelled', async () => { + let triggerCancel: (() => void) | undefined; + + const { job, publishedEvents, agentChatStreamingService, cancelCallbacks } = + buildJob({ + chatStream: createFakeChatStream({ + onFirstChunk: () => triggerCancel?.(), + }), + }); + + triggerCancel = () => cancelCallbacks.forEach((callback) => callback()); + + await job.handle(jobData); + + expect(publishedEvents.map((event) => event.type)).not.toContain( + 'stream-error', + ); + expect( + agentChatStreamingService.flushNextQueuedMessage, + ).not.toHaveBeenCalled(); + }); +});