feat: AI Workflow Builder agent (no-changelog) (#17423)

Co-authored-by: cubic-dev-ai[bot] <191113872+cubic-dev-ai[bot]@users.noreply.github.com>

packages/frontend/editor-ui/src/stores/builder.store.test.ts

@@ -1,7 +1,6 @@
-import { describe, it, expect, beforeEach } from 'vitest';
+import { describe, it, expect, beforeEach, vi } from 'vitest';
 import { setActivePinia, createPinia } from 'pinia';
 import { ENABLED_VIEWS, useBuilderStore } from '@/stores/builder.store';
-import type { ChatRequest } from '@/types/assistant.types';
 import { usePostHog } from './posthog.store';
 import { useSettingsStore } from '@/stores/settings.store';
 import { defaultSettings } from '../__tests__/defaults';
@@ -69,6 +68,12 @@ describe('AI Builder store', () => {
 		track.mockReset();
 	});
 
+	afterEach(() => {
+		vi.clearAllMocks();
+		vi.clearAllTimers();
+		vi.useRealTimers();
+	});
+
 	it('initializes with default values', () => {
 		const builderStore = useBuilderStore();
 
@@ -99,10 +104,10 @@ describe('AI Builder store', () => {
 		expect(builderStore.chatWidth).toBe(MAX_CHAT_WIDTH);
 	});
 
-	it('should open chat window', () => {
+	it('should open chat window', async () => {
 		const builderStore = useBuilderStore();
 
-		builderStore.openChat();
+		await builderStore.openChat();
 		expect(builderStore.chatWindowOpen).toBe(true);
 	});
 
@@ -113,101 +118,236 @@ describe('AI Builder store', () => {
 		expect(builderStore.chatWindowOpen).toBe(false);
 	});
 
-	it('can add a simple assistant message', () => {
+	it('can process a simple assistant message through API', async () => {
 		const builderStore = useBuilderStore();
 
-		const message: ChatRequest.MessageResponse = {
-			type: 'message',
-			role: 'assistant',
-			text: 'Hello!',
-		};
-		builderStore.addAssistantMessages([message], '1');
-		expect(builderStore.chatMessages.length).toBe(1);
-		expect(builderStore.chatMessages[0]).toEqual({
-			id: '1',
+		apiSpy.mockImplementationOnce((_ctx, _payload, onMessage, onDone) => {
+			onMessage({
+				messages: [
+					{
+						type: 'message',
+						role: 'assistant',
+						text: 'Hello!',
+					},
+				],
+				sessionId: 'test-session',
+			});
+			onDone();
+		});
+
+		builderStore.sendChatMessage({ text: 'Hi' });
+		await vi.waitFor(() => expect(builderStore.chatMessages.length).toBe(2));
+		expect(builderStore.chatMessages[0].role).toBe('user');
+		expect(builderStore.chatMessages[1]).toMatchObject({
 			type: 'text',
 			role: 'assistant',
 			content: 'Hello!',
-			quickReplies: undefined,
-			read: true, // Builder messages are always read
+			read: false,
 		});
 	});
 
-	it('can add a workflow step message', () => {
+	it('can process a workflow-updated message through API', async () => {
 		const builderStore = useBuilderStore();
 
-		const message: ChatRequest.MessageResponse = {
-			type: 'workflow-step',
-			role: 'assistant',
-			steps: ['Step 1', 'Step 2'],
-		};
-		builderStore.addAssistantMessages([message], '1');
-		expect(builderStore.chatMessages.length).toBe(1);
-		expect(builderStore.chatMessages[0]).toEqual({
-			id: '1',
-			type: 'workflow-step',
-			role: 'assistant',
-			steps: ['Step 1', 'Step 2'],
-			read: true,
+		apiSpy.mockImplementationOnce((_ctx, _payload, onMessage, onDone) => {
+			onMessage({
+				messages: [
+					{
+						type: 'workflow-updated',
+						role: 'assistant',
+						codeSnippet: '{"nodes":[],"connections":[]}',
+					},
+				],
+				sessionId: 'test-session',
+			});
+			onDone();
 		});
-	});
 
-	it('can add a workflow-generated message', () => {
-		const builderStore = useBuilderStore();
-
-		const message: ChatRequest.MessageResponse = {
-			type: 'workflow-generated',
+		builderStore.sendChatMessage({ text: 'Create workflow' });
+		await vi.waitFor(() => expect(builderStore.chatMessages.length).toBe(2));
+		expect(builderStore.chatMessages[1]).toMatchObject({
+			type: 'workflow-updated',
 			role: 'assistant',
 			codeSnippet: '{"nodes":[],"connections":[]}',
-		};
-		builderStore.addAssistantMessages([message], '1');
-		expect(builderStore.chatMessages.length).toBe(1);
-		expect(builderStore.chatMessages[0]).toEqual({
-			id: '1',
-			type: 'workflow-generated',
-			role: 'assistant',
-			codeSnippet: '{"nodes":[],"connections":[]}',
-			read: true,
+			read: false,
 		});
+
+		// Verify workflow messages are accessible via computed property
+		expect(builderStore.workflowMessages.length).toBe(1);
 	});
 
-	it('can add a rate-workflow message', () => {
+	it('should show processing results message when tools complete', async () => {
+		vi.useFakeTimers();
 		const builderStore = useBuilderStore();
 
-		const message: ChatRequest.MessageResponse = {
-			type: 'rate-workflow',
-			role: 'assistant',
-			content: 'How was the workflow?',
-		};
-		builderStore.addAssistantMessages([message], '1');
-		expect(builderStore.chatMessages.length).toBe(1);
-		expect(builderStore.chatMessages[0]).toEqual({
-			id: '1',
-			type: 'rate-workflow',
-			role: 'assistant',
-			content: 'How was the workflow?',
-			read: true,
+		// eslint-disable-next-line @typescript-eslint/no-explicit-any
+		let onMessageCallback: any;
+		// eslint-disable-next-line @typescript-eslint/no-explicit-any
+		let onDoneCallback: any;
+
+		apiSpy.mockImplementationOnce((_ctx, _payload, onMessage, onDone) => {
+			onMessageCallback = onMessage;
+			onDoneCallback = onDone;
 		});
-	});
 
-	it('should reset builder chat session', () => {
-		const builderStore = useBuilderStore();
+		builderStore.sendChatMessage({ text: 'Add nodes and connect them' });
 
-		const message: ChatRequest.MessageResponse = {
-			type: 'message',
-			role: 'assistant',
-			text: 'Hello!',
-			quickReplies: [
-				{ text: 'Yes', type: 'text' },
-				{ text: 'No', type: 'text' },
+		// Initially shows "Thinking..." from prepareForStreaming
+		expect(builderStore.assistantThinkingMessage).toBe('Thinking...');
+
+		// First tool starts
+		onMessageCallback({
+			messages: [
+				{
+					type: 'tool',
+					role: 'assistant',
+					toolName: 'add_nodes',
+					toolCallId: 'call-1',
+					status: 'running',
+					updates: [{ type: 'input', data: {} }],
+				},
 			],
-		};
-		builderStore.addAssistantMessages([message], '1');
-		expect(builderStore.chatMessages.length).toBe(1);
+		});
+
+		// Should show "Running tools..."
+		expect(builderStore.assistantThinkingMessage).toBe('Running tools...');
+
+		// Second tool starts (different toolCallId)
+		onMessageCallback({
+			messages: [
+				{
+					type: 'tool',
+					role: 'assistant',
+					toolName: 'connect_nodes',
+					toolCallId: 'call-2',
+					status: 'running',
+					updates: [{ type: 'input', data: {} }],
+				},
+			],
+		});
+
+		// Still showing "Running tools..." with multiple tools
+		expect(builderStore.assistantThinkingMessage).toBe('Running tools...');
+
+		// First tool completes
+		onMessageCallback({
+			messages: [
+				{
+					type: 'tool',
+					role: 'assistant',
+					toolName: 'add_nodes',
+					toolCallId: 'call-1',
+					status: 'completed',
+					updates: [{ type: 'output', data: { success: true } }],
+				},
+			],
+		});
+
+		// Still "Running tools..." because second tool is still running
+		expect(builderStore.assistantThinkingMessage).toBe('Running tools...');
+
+		// Second tool completes
+		onMessageCallback({
+			messages: [
+				{
+					type: 'tool',
+					role: 'assistant',
+					toolName: 'connect_nodes',
+					toolCallId: 'call-2',
+					status: 'completed',
+					updates: [{ type: 'output', data: { success: true } }],
+				},
+			],
+		});
+
+		// Now should show "Processing results..." because all tools completed
+		expect(builderStore.assistantThinkingMessage).toBe('Processing results...');
+
+		// Call onDone to stop streaming
+		onDoneCallback();
+
+		// Message should persist after streaming ends
+		expect(builderStore.streaming).toBe(false);
+		expect(builderStore.assistantThinkingMessage).toBe('Processing results...');
+
+		vi.useRealTimers();
+	});
+
+	it('should keep processing message when workflow-updated arrives', async () => {
+		const builderStore = useBuilderStore();
+
+		apiSpy.mockImplementationOnce((_ctx, _payload, onMessage, onDone) => {
+			// Tool completes
+			onMessage({
+				messages: [
+					{
+						type: 'tool',
+						role: 'assistant',
+						toolName: 'add_nodes',
+						toolCallId: 'call-1',
+						status: 'completed',
+						updates: [{ type: 'output', data: { success: true } }],
+					},
+				],
+			});
+
+			// Workflow update arrives
+			onMessage({
+				messages: [
+					{
+						type: 'workflow-updated',
+						role: 'assistant',
+						codeSnippet: '{"nodes": [], "connections": {}}',
+					},
+				],
+			});
+
+			// Call onDone to stop streaming
+			onDone();
+		});
+
+		builderStore.sendChatMessage({ text: 'Add a node' });
+
+		// Should show "Processing results..." when tool completes
+		await vi.waitFor(() =>
+			expect(builderStore.assistantThinkingMessage).toBe('Processing results...'),
+		);
+
+		// Should still show "Processing results..." after workflow-updated
+		await vi.waitFor(() => expect(builderStore.chatMessages).toHaveLength(3)); // user + tool + workflow
+		expect(builderStore.assistantThinkingMessage).toBe('Processing results...');
+
+		// Verify streaming has ended
+		expect(builderStore.streaming).toBe(false);
+	});
+
+	it('should reset builder chat session', async () => {
+		const builderStore = useBuilderStore();
+
+		apiSpy.mockImplementationOnce((_ctx, _payload, onMessage, onDone) => {
+			onMessage({
+				messages: [
+					{
+						type: 'message',
+						role: 'assistant',
+						text: 'Hello!',
+						quickReplies: [
+							{ text: 'Yes', type: 'text' },
+							{ text: 'No', type: 'text' },
+						],
+					},
+				],
+				sessionId: 'test-session',
+			});
+			onDone();
+		});
+
+		builderStore.sendChatMessage({ text: 'Hi' });
+		await vi.waitFor(() => expect(builderStore.chatMessages.length).toBe(2));
 
 		builderStore.resetBuilderChat();
 		expect(builderStore.chatMessages).toEqual([]);
-		expect(builderStore.currentSessionId).toBeUndefined();
+		expect(builderStore.assistantThinkingMessage).toBeUndefined();
 	});
 
 	it('should not show builder if disabled in settings', () => {
@@ -259,13 +399,13 @@ describe('AI Builder store', () => {
 			onDone();
 		});
 
-		await builderStore.initBuilderChat('I want to build a workflow', 'chat');
+		builderStore.sendChatMessage({ text: 'I want to build a workflow' });
+		await vi.waitFor(() => expect(builderStore.chatMessages.length).toBe(2));
 
 		expect(apiSpy).toHaveBeenCalled();
-		expect(builderStore.currentSessionId).toEqual(mockSessionId);
-		expect(builderStore.chatMessages.length).toBe(2); // user message + assistant response
 		expect(builderStore.chatMessages[0].role).toBe('user');
 		expect(builderStore.chatMessages[1].role).toBe('assistant');
+		expect(builderStore.streaming).toBe(false);
 	});
 
 	it('should send a follow-up message in an existing session', async () => {
@@ -302,18 +442,15 @@ describe('AI Builder store', () => {
 			onDone();
 		});
 
-		await builderStore.initBuilderChat('I want to build a workflow', 'chat');
-
-		// Should be 2 messages now (user question + assistant response)
-		expect(builderStore.chatMessages.length).toBe(2);
+		builderStore.sendChatMessage({ text: 'I want to build a workflow' });
+		await vi.waitFor(() => expect(builderStore.chatMessages.length).toBe(2));
 
 		// Send a follow-up message
-		await builderStore.sendMessage({ text: 'Generate a workflow for me' });
+		builderStore.sendChatMessage({ text: 'Generate a workflow for me' });
+		await vi.waitFor(() => expect(builderStore.chatMessages.length).toBe(4));
 
 		const thirdMessage = builderStore.chatMessages[2] as ChatUI.TextMessage;
 		const fourthMessage = builderStore.chatMessages[3] as ChatUI.TextMessage;
-		// Should be 4 messages now (2 initial + user follow-up + assistant response)
-		expect(builderStore.chatMessages.length).toBe(4);
 		expect(thirdMessage.role).toBe('user');
 		expect(thirdMessage.type).toBe('text');
 		expect(thirdMessage.content).toBe('Generate a workflow for me');
@@ -330,10 +467,8 @@ describe('AI Builder store', () => {
 			onError(new Error('An API error occurred'));
 		});
 
-		await builderStore.initBuilderChat('I want to build a workflow', 'chat');
-
-		// Should have user message + error message
-		expect(builderStore.chatMessages.length).toBe(2);
+		builderStore.sendChatMessage({ text: 'I want to build a workflow' });
+		await vi.waitFor(() => expect(builderStore.chatMessages.length).toBe(2));
 		expect(builderStore.chatMessages[0].role).toBe('user');
 		expect(builderStore.chatMessages[1].type).toBe('error');
 
@@ -341,6 +476,9 @@ describe('AI Builder store', () => {
 		const errorMessage = builderStore.chatMessages[1] as ChatUI.ErrorMessage;
 		expect(errorMessage.retry).toBeDefined();
 
+		// Verify streaming state was reset
+		expect(builderStore.streaming).toBe(false);
+
 		// Set up a successful response for the retry
 		apiSpy.mockImplementationOnce((_ctx, _payload, onMessage, onDone) => {
 			onMessage({
@@ -357,10 +495,10 @@ describe('AI Builder store', () => {
 		});
 
 		// Retry the failed request
-		await errorMessage.retry?.();
-
-		// Should now have just the user message and success message
-		expect(builderStore.chatMessages.length).toBe(2);
+		if (errorMessage.retry) {
+			void errorMessage.retry();
+			await vi.waitFor(() => expect(builderStore.chatMessages.length).toBe(2));
+		}
 		expect(builderStore.chatMessages[0].role).toBe('user');
 		expect(builderStore.chatMessages[1].type).toBe('text');
 		expect((builderStore.chatMessages[1] as ChatUI.TextMessage).content).toBe(