feat: Abort AI builder requests on chat stop (#17854)

packages/cli/src/controllers/__tests__/ai.controller.test.ts

@@ -152,6 +152,7 @@ describe('AiController', () => {
 					},
 				},
 				request.user,
+				expect.any(AbortSignal),
 			);
 			expect(response.header).toHaveBeenCalledWith('Content-type', 'application/json-lines');
 			expect(response.flush).toHaveBeenCalled();
@@ -241,5 +242,157 @@ describe('AiController', () => {
 			expect(response.json).not.toHaveBeenCalled();
 			expect(response.end).toHaveBeenCalled();
 		});
+
+		describe('Abort handling', () => {
+			it('should create AbortController and handle connection close', async () => {
+				let abortHandler: (() => void) | undefined;
+				let abortSignalPassed: AbortSignal | undefined;
+
+				// Mock response.on to capture the close handler
+				response.on.mockImplementation((event: string, handler: () => void) => {
+					if (event === 'close') {
+						abortHandler = handler;
+					}
+					return response;
+				});
+
+				// Create a generator that yields once then checks for abort
+				async function* testGenerator() {
+					yield {
+						messages: [{ role: 'assistant', type: 'message', text: 'Processing...' } as const],
+					};
+					// Check if aborted and throw if so
+					if (abortSignalPassed?.aborted) {
+						throw new Error('Aborted');
+					}
+				}
+
+				workflowBuilderService.chat.mockImplementation((_payload, _user, signal) => {
+					abortSignalPassed = signal;
+					return testGenerator();
+				});
+
+				// Start the request (but don't await it)
+				const buildPromise = controller.build(request, response, payload);
+
+				// Wait a bit to ensure the generator is created and starts processing
+				await new Promise((resolve) => setTimeout(resolve, 50));
+
+				// Verify abort signal was passed to the service
+				expect(abortSignalPassed).toBeDefined();
+				expect(abortSignalPassed).toBeInstanceOf(AbortSignal);
+				expect(abortSignalPassed?.aborted).toBe(false);
+
+				// Verify close handler was registered
+				expect(response.on).toHaveBeenCalledWith('close', expect.any(Function));
+				expect(abortHandler).toBeDefined();
+
+				// Simulate connection close
+				abortHandler!();
+
+				// Verify the signal was aborted
+				expect(abortSignalPassed?.aborted).toBe(true);
+
+				// Wait for the promise to settle
+				await buildPromise.catch(() => {
+					// Expected to throw due to abort
+				});
+
+				// Verify response was ended
+				expect(response.end).toHaveBeenCalled();
+			});
+
+			it('should pass abort signal to workflow builder service', async () => {
+				let capturedSignal: AbortSignal | undefined;
+
+				async function* mockGenerator() {
+					yield { messages: [{ role: 'assistant', type: 'message', text: 'Test' } as const] };
+				}
+
+				workflowBuilderService.chat.mockImplementation((_payload, _user, signal) => {
+					capturedSignal = signal;
+					return mockGenerator();
+				});
+
+				await controller.build(request, response, payload);
+
+				expect(capturedSignal).toBeDefined();
+				expect(capturedSignal).toBeInstanceOf(AbortSignal);
+				expect(workflowBuilderService.chat).toHaveBeenCalledWith(
+					expect.any(Object),
+					request.user,
+					capturedSignal,
+				);
+			});
+
+			it('should handle stream interruption when connection closes', async () => {
+				let abortHandler: (() => void) | undefined;
+				let abortSignalPassed: AbortSignal | undefined;
+
+				response.on.mockImplementation((event: string, handler: () => void) => {
+					if (event === 'close') {
+						abortHandler = handler;
+					}
+					return response;
+				});
+
+				// Create a generator that yields multiple chunks
+				async function* mockChatGenerator() {
+					yield { messages: [{ role: 'assistant', type: 'message', text: 'Chunk 1' } as const] };
+
+					// Check if aborted before yielding next chunk
+					if (abortSignalPassed?.aborted) {
+						throw new Error('Aborted');
+					}
+
+					// This second chunk should not be reached if aborted
+					yield { messages: [{ role: 'assistant', type: 'message', text: 'Chunk 2' } as const] };
+				}
+
+				workflowBuilderService.chat.mockImplementation((_payload, _user, signal) => {
+					abortSignalPassed = signal;
+					return mockChatGenerator();
+				});
+
+				// Start the build process
+				const buildPromise = controller.build(request, response, payload);
+
+				// Wait for first chunk to be written
+				await new Promise((resolve) => setTimeout(resolve, 20));
+
+				// Should have written at least one chunk
+				expect(response.write).toHaveBeenCalled();
+				const writeCallsBeforeAbort = response.write.mock.calls.length;
+
+				// Simulate connection close
+				abortHandler!();
+
+				// Wait for the build to complete
+				await buildPromise.catch(() => {
+					// Expected to catch abort error
+				});
+
+				// Should not have written additional chunks after abort
+				expect(response.write).toHaveBeenCalledTimes(writeCallsBeforeAbort);
+				expect(response.end).toHaveBeenCalled();
+			});
+
+			it('should cleanup abort listener on successful completion', async () => {
+				const onSpy = jest.spyOn(response, 'on');
+				const offSpy = jest.spyOn(response, 'off');
+
+				async function* mockGenerator() {
+					yield { messages: [{ role: 'assistant', type: 'message', text: 'Complete' } as const] };
+				}
+
+				workflowBuilderService.chat.mockReturnValue(mockGenerator());
+
+				await controller.build(request, response, payload);
+
+				// Verify close handler was registered and then removed
+				expect(onSpy).toHaveBeenCalledWith('close', expect.any(Function));
+				expect(offSpy).toHaveBeenCalledWith('close', expect.any(Function));
+			});
+		});
 	});
 });

packages/cli/src/controllers/ai.controller.ts

@@ -46,6 +46,13 @@ export class AiController {
 		@Body payload: AiBuilderChatRequestDto,
 	) {
 		try {
+			const abortController = new AbortController();
+			const { signal } = abortController;
+
+			const handleClose = () => abortController.abort();
+
+			res.on('close', handleClose);
+
 			const { text, workflowContext } = payload.payload;
 			const aiResponse = this.workflowBuilderService.chat(
 				{
@@ -57,6 +64,7 @@ export class AiController {
 					},
 				},
 				req.user,
+				signal,
 			);
 
 			res.header('Content-type', 'application/json-lines').flush();
@@ -83,6 +91,9 @@ export class AiController {
 					],
 				};
 				res.write(JSON.stringify(errorChunk) + '⧉⇋⇋➽⌑⧉§§\n');
+			} finally {
+				// Clean up event listener
+				res.off('close', handleClose);
 			}
 
 			res.end();

packages/cli/src/services/ai-workflow-builder.service.ts

@@ -48,9 +48,9 @@ export class WorkflowBuilderService {
 		return this.service;
 	}
 
-	async *chat(payload: ChatPayload, user: IUser) {
+	async *chat(payload: ChatPayload, user: IUser, abortSignal?: AbortSignal) {
 		const service = await this.getService();
-		yield* service.chat(payload, user);
+		yield* service.chat(payload, user, abortSignal);
 	}
 
 	async getSessions(workflowId: string | undefined, user: IUser) {