test(ai): pin StreamAgentChatJob stream lifecycle with a reusable spec harness (#22479)
## 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)_ <!-- This is an auto-generated description by cubic. --> <a href="https://cubic.dev/pr/twentyhq/twenty/pull/22479?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> <!-- End of auto-generated description by cubic. -->
This commit is contained in:
Félix Malfait
+318
@@ -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<string, unknown>;
|
||||
|
||||
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> | void;
|
||||
}) =>
|
||||
new ReadableStream<UIMessageChunk>({
|
||||
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<typeof createFakeChatStream>;
|
||||
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();
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user