feat(core): Make multiple replies to a single message work (no-changelog) (#21211)

packages/@n8n/api-types/src/chat-hub.ts

@@ -1,3 +1,4 @@
+import type { StructuredChunk } from 'n8n-workflow';
 import { z } from 'zod';
 import { Z } from 'zod-class';
 
@@ -94,7 +95,6 @@ export type ChatModelsResponse = Record<
 export class ChatHubSendMessageRequest extends Z.class({
 	messageId: z.string().uuid(),
 	sessionId: z.string().uuid(),
-	replyId: z.string().uuid(),
 	message: z.string(),
 	model: chatHubConversationModelSchema,
 	previousMessageId: z.string().uuid().nullable(),
@@ -107,7 +107,6 @@ export class ChatHubSendMessageRequest extends Z.class({
 }) {}
 
 export class ChatHubRegenerateMessageRequest extends Z.class({
-	replyId: z.string().uuid(),
 	model: chatHubConversationModelSchema,
 	credentials: z.record(
 		z.object({
@@ -120,7 +119,6 @@ export class ChatHubRegenerateMessageRequest extends Z.class({
 export class ChatHubEditMessageRequest extends Z.class({
 	message: z.string(),
 	messageId: z.string().uuid(),
-	replyId: z.string().uuid(),
 	model: chatHubConversationModelSchema,
 	credentials: z.record(
 		z.object({
@@ -222,3 +220,11 @@ export class ChatHubUpdateAgentRequest extends Z.class({
 	provider: chatHubProviderSchema.optional(),
 	model: z.string().max(64).optional(),
 }) {}
+
+export interface EnrichedStructuredChunk extends StructuredChunk {
+	metadata: StructuredChunk['metadata'] & {
+		messageId: ChatMessageId;
+		previousMessageId: ChatMessageId | null;
+		retryOfMessageId: ChatMessageId | null;
+	};
+}

packages/@n8n/api-types/src/index.ts

@@ -32,6 +32,7 @@ export {
 	type ChatHubAgentDto,
 	ChatHubCreateAgentRequest,
 	ChatHubUpdateAgentRequest,
+	type EnrichedStructuredChunk,
 } from './chat-hub';
 
 export type { Collaborator } from './push/collaboration';

packages/cli/src/modules/chat-hub/chat-hub.controller.ts

@@ -26,12 +26,12 @@ import {
 import type { Response } from 'express';
 import { strict as assert } from 'node:assert';
 
+import { ResponseError } from '@/errors/response-errors/abstract/response.error';
+
 import { ChatHubService } from './chat-hub.service';
 import { ChatHubAgentService } from './chat-hub-agent.service';
 import { ChatModelsRequestDto } from './dto/chat-models-request.dto';
 
-import { ResponseError } from '@/errors/response-errors/abstract/response.error';
-
 @RestController('/chat')
 export class ChatHubController {
 	constructor(

packages/cli/src/modules/chat-hub/chat-hub.service.ts

@@ -11,6 +11,7 @@ import {
 	ChatHubConversationModel,
 	ChatHubMessageStatus,
 	chatHubProviderSchema,
+	type EnrichedStructuredChunk,
 } from '@n8n/api-types';
 import { Logger } from '@n8n/backend-common';
 import {
@@ -46,16 +47,8 @@ import {
 } from 'n8n-workflow';
 import { v4 as uuidv4 } from 'uuid';
 
-import { ActiveExecutions } from '@/active-executions';
-import { BadRequestError } from '@/errors/response-errors/bad-request.error';
-import { NotFoundError } from '@/errors/response-errors/not-found.error';
-import { ExecutionService } from '@/executions/execution.service';
-import { DynamicNodeParametersService } from '@/services/dynamic-node-parameters.service';
-import { getBase } from '@/workflow-execute-additional-data';
-import { WorkflowExecutionService } from '@/workflows/workflow-execution.service';
-import { WorkflowFinderService } from '@/workflows/workflow-finder.service';
-import { WorkflowService } from '@/workflows/workflow.service';
-
+import { ChatHubAgentService } from './chat-hub-agent.service';
+import { ChatHubCredentialsService } from './chat-hub-credentials.service';
 import type { ChatHubMessage } from './chat-hub-message.entity';
 import { CONVERSATION_TITLE_GENERATION_PROMPT } from './chat-hub.constants';
 import type {
@@ -68,10 +61,18 @@ import type {
 import { ChatHubMessageRepository } from './chat-message.repository';
 import { ChatHubSessionRepository } from './chat-session.repository';
 import { getMaxContextWindowTokens } from './context-limits';
-import { captureResponseWrites } from './stream-capturer';
-import { ChatHubAgentService } from './chat-hub-agent.service';
-import { ChatHubCredentialsService } from './chat-hub-credentials.service';
+import { interceptResponseWrites, createStructuredChunkAggregator } from './stream-capturer';
+
+import { ActiveExecutions } from '@/active-executions';
 import { CredentialsFinderService } from '@/credentials/credentials-finder.service';
+import { BadRequestError } from '@/errors/response-errors/bad-request.error';
+import { NotFoundError } from '@/errors/response-errors/not-found.error';
+import { ExecutionService } from '@/executions/execution.service';
+import { DynamicNodeParametersService } from '@/services/dynamic-node-parameters.service';
+import { getBase } from '@/workflow-execute-additional-data';
+import { WorkflowExecutionService } from '@/workflows/workflow-execution.service';
+import { WorkflowFinderService } from '@/workflows/workflow-finder.service';
+import { WorkflowService } from '@/workflows/workflow.service';
 
 const providerNodeTypeMapping: Record<ChatHubLLMProvider, INodeTypeNameVersion> = {
 	openai: {
@@ -132,7 +133,7 @@ export class ChatHubService {
 		credentialIds: Record<ChatHubLLMProvider, string | null>,
 	): Promise<ChatModelsResponse> {
 		const additionalData = await getBase({ userId: user.id });
-		const providers = chatHubProviderSchema.options.filter((provider) => provider !== 'n8n');
+		const providers = chatHubProviderSchema.options;
 
 		const allCredentials = await this.credentialsFinderService.findCredentialsForUser(user, [
 			'credential:read',
@@ -439,28 +440,24 @@ export class ChatHubService {
 	}
 
 	async sendHumanMessage(res: Response, user: User, payload: HumanMessagePayload) {
-		const { sessionId, messageId, replyId, message } = payload;
+		const { sessionId, messageId, message, model, credentials, previousMessageId } = payload;
 		const provider = payload.model.provider;
 
-		const selectedModel: ModelWithCredentials = this.getModelWithCredentials(
-			payload.model,
-			payload.credentials,
-		);
+		const selectedModel: ModelWithCredentials = {
+			...payload.model,
+			credentialId: provider !== 'n8n' ? this.pickCredentialId(provider, credentials) : null,
+		};
 
 		const workflow = await this.messageRepository.manager.transaction(async (trx) => {
 			const session = await this.getChatSession(user, sessionId, selectedModel, true, trx);
-			await this.ensurePreviousMessage(payload.previousMessageId, sessionId, trx);
+			await this.ensurePreviousMessage(previousMessageId, sessionId, trx);
 			const messages = Object.fromEntries((session.messages ?? []).map((m) => [m.id, m]));
-			const history = this.buildMessageHistory(messages, payload.previousMessageId);
+			const history = this.buildMessageHistory(messages, previousMessageId);
 
-			await this.saveHumanMessage(
-				payload,
-				user,
-				payload.previousMessageId,
-				selectedModel,
-				undefined,
-				trx,
-			);
+			await this.saveHumanMessage(payload, user, previousMessageId, selectedModel, undefined, trx);
+
+			// generate title on receiving the first human message only
+			const generateTitle = previousMessageId === null;
 
 			if (provider === 'n8n') {
 				return await this.prepareCustomAgentWorkflow(
@@ -475,53 +472,235 @@ export class ChatHubService {
 				return await this.prepareChatAgentWorkflow(
 					payload.model.agentId,
 					user,
-					payload,
 					sessionId,
 					history,
 					message,
+					generateTitle,
 					trx,
 				);
 			}
 
-			return await this.prepareBaseChatWorkflow(user, payload, sessionId, history, message, trx);
+			return await this.prepareBaseChatWorkflow(
+				user,
+				sessionId,
+				credentials,
+				model,
+				history,
+				message,
+				generateTitle,
+				trx,
+			);
 		});
 
-		try {
-			await this.executeChatWorkflow(
-				res,
-				user,
-				workflow,
-				replyId,
-				sessionId,
-				messageId,
-				selectedModel,
-			);
-		} finally {
-			if (provider !== 'n8n') {
-				await this.deleteChatWorkflow(workflow.workflowData.id);
+		await this.executeChatWorkflowWithCleanup(
+			res,
+			user,
+			workflow,
+			sessionId,
+			messageId,
+			selectedModel,
+			provider,
+		);
+	}
+
+	async editMessage(res: Response, user: User, payload: EditMessagePayload) {
+		const { sessionId, editId, messageId, message, model, credentials } = payload;
+		const provider = payload.model.provider;
+
+		const selectedModel: ModelWithCredentials = this.getModelWithCredentials(model, credentials);
+
+		const workflow = await this.messageRepository.manager.transaction(async (trx) => {
+			const session = await this.getChatSession(user, sessionId, selectedModel, true, trx);
+			const messageToEdit = await this.getChatMessage(session.id, editId, [], trx);
+
+			if (!['ai', 'human'].includes(messageToEdit.type)) {
+				throw new BadRequestError('Only human and AI messages can be edited');
 			}
+
+			if (messageToEdit.type === 'ai') {
+				// AI edits just change the original message without revisioning or response generation
+				await this.messageRepository.updateChatMessage(editId, { content: payload.message }, trx);
+				return null;
+			}
+
+			if (messageToEdit.type === 'human') {
+				const messages = Object.fromEntries((session.messages ?? []).map((m) => [m.id, m]));
+				const history = this.buildMessageHistory(messages, messageToEdit.previousMessageId);
+
+				// If the message to edit isn't the original message, we want to point to the original message
+				const revisionOfMessageId = messageToEdit.revisionOfMessageId ?? messageToEdit.id;
+
+				await this.saveHumanMessage(
+					payload,
+					user,
+					messageToEdit.previousMessageId,
+					selectedModel,
+					revisionOfMessageId,
+					trx,
+				);
+
+				if (provider === 'n8n') {
+					return await this.prepareCustomAgentWorkflow(
+						user,
+						sessionId,
+						payload.model.workflowId,
+						message,
+					);
+				}
+
+				if (provider === 'custom-agent') {
+					return await this.prepareChatAgentWorkflow(
+						payload.model.agentId,
+						user,
+						sessionId,
+						history,
+						message,
+						false,
+						trx,
+					);
+				}
+
+				return await this.prepareBaseChatWorkflow(
+					user,
+					sessionId,
+					credentials,
+					model,
+					history,
+					message,
+					false,
+					trx,
+				);
+			}
+			return null;
+		});
+
+		if (!workflow) {
+			return;
 		}
+
+		await this.executeChatWorkflowWithCleanup(
+			res,
+			user,
+			workflow,
+			sessionId,
+			messageId,
+			selectedModel,
+			provider,
+		);
+	}
+
+	async regenerateAIMessage(res: Response, user: User, payload: RegenerateMessagePayload) {
+		const { sessionId, retryId, model, credentials } = payload;
+		const provider = payload.model.provider;
+
+		const selectedModel: ModelWithCredentials = {
+			...payload.model,
+			credentialId:
+				payload.model.provider !== 'n8n'
+					? this.pickCredentialId(payload.model.provider, payload.credentials)
+					: null,
+		};
+
+		const { workflow, retryOfMessageId, previousMessageId } =
+			await this.messageRepository.manager.transaction(async (trx) => {
+				const session = await this.getChatSession(user, sessionId, undefined, false, trx);
+				const messageToRetry = await this.getChatMessage(session.id, retryId, [], trx);
+
+				if (messageToRetry.type !== 'ai') {
+					throw new BadRequestError('Can only retry AI messages');
+				}
+
+				const messages = Object.fromEntries((session.messages ?? []).map((m) => [m.id, m]));
+				const history = this.buildMessageHistory(messages, messageToRetry.previousMessageId);
+
+				const lastHumanMessage = history.filter((m) => m.type === 'human').pop();
+				if (!lastHumanMessage) {
+					throw new BadRequestError('No human message found to base the retry on');
+				}
+
+				// Remove any (AI) messages that came after the last human message
+				const lastHumanMessageIndex = history.indexOf(lastHumanMessage);
+				if (lastHumanMessageIndex !== -1) {
+					history.splice(lastHumanMessageIndex + 1);
+				}
+
+				// Rerun the workflow, replaying the last human message
+
+				// If the message being retried is itself a retry, we want to point to the original message
+				const retryOfMessageId = messageToRetry.retryOfMessageId ?? messageToRetry.id;
+				const message = lastHumanMessage ? lastHumanMessage.content : '';
+
+				let workflow;
+				if (provider === 'n8n') {
+					workflow = await this.prepareCustomAgentWorkflow(
+						user,
+						sessionId,
+						payload.model.workflowId,
+						message,
+					);
+				} else if (provider === 'custom-agent') {
+					workflow = await this.prepareChatAgentWorkflow(
+						payload.model.agentId,
+						user,
+						sessionId,
+						history,
+						message,
+						false,
+						trx,
+					);
+				} else {
+					workflow = await this.prepareBaseChatWorkflow(
+						user,
+						sessionId,
+						credentials,
+						model,
+						history,
+						message,
+						false,
+						trx,
+					);
+				}
+
+				return {
+					workflow,
+					previousMessageId: lastHumanMessage.id,
+					retryOfMessageId,
+				};
+			});
+
+		await this.executeChatWorkflowWithCleanup(
+			res,
+			user,
+			workflow,
+			sessionId,
+			previousMessageId,
+			selectedModel,
+			provider,
+			retryOfMessageId,
+		);
 	}
 
 	private async prepareBaseChatWorkflow(
 		user: User,
-		payload: HumanMessagePayload,
 		sessionId: ChatSessionId,
+		credentials: INodeCredentials,
+		model: ChatHubConversationModel,
 		history: ChatHubMessage[],
 		message: string,
+		generateConversationTitle: boolean,
 		trx: EntityManager,
 		systemMessage?: string,
 	) {
-		const credential = await this.ensureCredentials(user, payload.model, payload.credentials, trx);
+		const credential = await this.ensureCredentials(user, model, credentials, trx);
 
 		return await this.createChatWorkflow(
 			sessionId,
 			credential.projectId,
 			history,
 			message,
-			payload.credentials,
-			payload.model,
-			payload.previousMessageId === null, // generate title on receiving the first human message only
+			credentials,
+			model,
+			generateConversationTitle,
 			trx,
 			systemMessage,
 		);
@@ -530,10 +709,10 @@ export class ChatHubService {
 	private async prepareChatAgentWorkflow(
 		agentId: string,
 		user: User,
-		payload: HumanMessagePayload,
 		sessionId: ChatSessionId,
 		history: ChatHubMessage[],
 		message: string,
+		generateConversationTitle: boolean,
 		trx: EntityManager,
 	) {
 		const agent = await this.chatHubAgentService.getAgentById(agentId, user.id);
@@ -557,27 +736,27 @@ export class ChatHubService {
 
 		const systemMessage = agent.systemPrompt;
 
-		const updatedPayload: HumanMessagePayload = {
-			...payload,
-			model: {
-				provider: agent.provider,
-				model: agent.model,
-				name: agent.model,
-			} as ChatHubConversationModel,
-			credentials: {
-				[PROVIDER_CREDENTIAL_TYPE_MAP[agent.provider]]: {
-					id: credentialId,
-					name: '',
-				},
+		const model: ChatHubConversationModel = {
+			provider: agent.provider,
+			model: agent.model,
+			name: agent.model,
+		};
+
+		const credentials: INodeCredentials = {
+			[PROVIDER_CREDENTIAL_TYPE_MAP[agent.provider]]: {
+				id: credentialId,
+				name: '',
 			},
 		};
 
 		return await this.prepareBaseChatWorkflow(
 			user,
-			updatedPayload,
 			sessionId,
+			credentials,
+			model,
 			history,
 			message,
+			generateConversationTitle,
 			trx,
 			systemMessage,
 		);
@@ -618,10 +797,10 @@ export class ChatHubService {
 			);
 		}
 
-		const agents = workflowEntity.nodes.filter((node) => node.type === AGENT_LANGCHAIN_NODE_TYPE);
-		if (agents.length !== 1) {
-			throw new BadRequestError('Workflow must have exactly one AI Agent node');
-		}
+		// const agents = workflowEntity.nodes.filter((node) => node.type === AGENT_LANGCHAIN_NODE_TYPE);
+		// if (agents.length !== 1) {
+		// 	throw new BadRequestError('Workflow must have exactly one AI Agent node');
+		// }
 
 		return {
 			workflowData: {
@@ -677,160 +856,6 @@ export class ChatHubService {
 		}
 	}
 
-	async editMessage(res: Response, user: User, payload: EditMessagePayload) {
-		const { sessionId, editId, messageId, replyId } = payload;
-
-		const selectedModel: ModelWithCredentials = this.getModelWithCredentials(
-			payload.model,
-			payload.credentials,
-		);
-
-		const workflow = await this.messageRepository.manager.transaction(async (trx) => {
-			const credential = await this.ensureCredentials(
-				user,
-				payload.model,
-				payload.credentials,
-				trx,
-			);
-			const session = await this.getChatSession(user, sessionId, undefined, false, trx);
-			const messageToEdit = await this.getChatMessage(session.id, editId, [], trx);
-
-			if (!['ai', 'human'].includes(messageToEdit.type)) {
-				throw new BadRequestError('Only human and AI messages can be edited');
-			}
-
-			if (messageToEdit.type === 'ai') {
-				// AI edits just change the original message without revisioning or response generation
-				await this.messageRepository.updateChatMessage(editId, { content: payload.message }, trx);
-				return null;
-			}
-
-			if (messageToEdit.type === 'human') {
-				// Human messages branch the conversation and trigger an AI reply
-				const { message } = payload;
-				const messages = Object.fromEntries((session.messages ?? []).map((m) => [m.id, m]));
-				const history = this.buildMessageHistory(messages, messageToEdit.previousMessageId);
-
-				// If the message to edit isn't the original message, we want to point to the original message
-				const revisionOfMessageId = messageToEdit.revisionOfMessageId ?? messageToEdit.id;
-
-				await this.saveHumanMessage(
-					payload,
-					user,
-					messageToEdit.previousMessageId,
-					selectedModel,
-					revisionOfMessageId,
-					trx,
-				);
-
-				return await this.createChatWorkflow(
-					session.id,
-					credential.projectId,
-					history,
-					message,
-					payload.credentials,
-					payload.model,
-					messageToEdit.previousMessageId === null,
-					trx,
-				);
-			}
-
-			return null;
-		});
-
-		if (!workflow) {
-			return;
-		}
-
-		try {
-			await this.executeChatWorkflow(
-				res,
-				user,
-				workflow,
-				replyId,
-				sessionId,
-				messageId,
-				selectedModel,
-			);
-		} finally {
-			await this.deleteChatWorkflow(workflow.workflowData.id);
-		}
-	}
-
-	async regenerateAIMessage(res: Response, user: User, payload: RegenerateMessagePayload) {
-		const { sessionId, retryId, replyId } = payload;
-		const selectedModel: ModelWithCredentials = this.getModelWithCredentials(
-			payload.model,
-			payload.credentials,
-		);
-
-		const { workflow, retryOfMessageId, previousMessageId } =
-			await this.messageRepository.manager.transaction(async (trx) => {
-				const credential = await this.ensureCredentials(
-					user,
-					payload.model,
-					payload.credentials,
-					trx,
-				);
-				const session = await this.getChatSession(user, sessionId, undefined, false, trx);
-				const messageToRetry = await this.getChatMessage(session.id, retryId, [], trx);
-
-				if (messageToRetry.type !== 'ai') {
-					throw new BadRequestError('Can only retry AI messages');
-				}
-
-				const messages = Object.fromEntries((session.messages ?? []).map((m) => [m.id, m]));
-				const history = this.buildMessageHistory(messages, messageToRetry.previousMessageId);
-
-				const lastHumanMessage = history.filter((m) => m.type === 'human').pop();
-				if (!lastHumanMessage) {
-					throw new BadRequestError('No human message found to base the retry on');
-				}
-
-				// Remove any (AI) messages that came after the last human message
-				const lastHumanMessageIndex = history.indexOf(lastHumanMessage);
-				if (lastHumanMessageIndex !== -1) {
-					history.splice(lastHumanMessageIndex + 1);
-				}
-
-				// Rerun the workflow, replaying the last human message
-
-				// If the message being retried is itself a retry, we want to point to the original message
-				const retryOfMessageId = messageToRetry.retryOfMessageId ?? messageToRetry.id;
-				const workflow = await this.createChatWorkflow(
-					session.id,
-					credential.projectId,
-					history,
-					lastHumanMessage ? lastHumanMessage.content : '',
-					payload.credentials,
-					payload.model,
-					false,
-					trx,
-				);
-
-				return {
-					workflow,
-					previousMessageId: lastHumanMessage.id,
-					retryOfMessageId,
-				};
-			});
-
-		try {
-			await this.executeChatWorkflow(
-				res,
-				user,
-				workflow,
-				replyId,
-				sessionId,
-				previousMessageId,
-				selectedModel,
-				retryOfMessageId,
-			);
-		} finally {
-			await this.deleteChatWorkflow(workflow.workflowData.id);
-		}
-	}
-
 	async stopGeneration(user: User, sessionId: ChatSessionId, messageId: ChatMessageId) {
 		const session = await this.getChatSession(user, sessionId);
 		const message = await this.getChatMessage(session.id, messageId, [
@@ -861,11 +886,10 @@ export class ChatHubService {
 			workflowData: IWorkflowBase;
 			triggerToStartFrom: { name: string; data?: ITaskData };
 		},
-		replyId: ChatMessageId,
 		sessionId: ChatSessionId,
 		previousMessageId: ChatMessageId,
 		selectedModel: ModelWithCredentials,
-		retryOfMessageId?: ChatMessageId,
+		retryOfMessageId: ChatMessageId | null = null,
 	) {
 		const { workflowData, triggerToStartFrom } = workflow;
 
@@ -875,19 +899,96 @@ export class ChatHubService {
 
 		// Capture the streaming response as it's being generated to save
 		// partial messages in the database when generation gets cancelled.
-		let partialMessage = '';
-		const onChunk = (chunk: string) => {
-			const data = jsonParse<StructuredChunk>(chunk);
-			if (data && data.type === 'item' && typeof data.content === 'string') {
-				partialMessage += data.content;
+		let executionId: string | undefined = undefined;
+
+		const aggregator = createStructuredChunkAggregator(previousMessageId, retryOfMessageId, {
+			onBegin: async (message) => {
+				await this.saveAIMessage({
+					...message,
+					sessionId,
+					executionId,
+					selectedModel,
+					retryOfMessageId,
+				});
+			},
+			onItem: (_message, _chunk) => {
+				// We could save partial messages to DB here if we wanted to,
+				// but they would be very frequent updates.
+			},
+			onEnd: async (message) => {
+				await this.messageRepository.updateChatMessage(message.id, {
+					content: message.content,
+					status: message.status,
+				});
+			},
+			onError: async (message, _errorText) => {
+				await this.messageRepository.manager.transaction(async (trx) => {
+					// Always update the content to whatever was generated so far, including the possible error text
+					await this.messageRepository.updateChatMessage(
+						message.id,
+						{
+							content: message.content,
+						},
+						trx,
+					);
+
+					// When messages are cancelled they're already marked cancelled on `stopGeneration`
+					const savedMessage = await this.messageRepository.getOneById(
+						message.id,
+						sessionId,
+						[],
+						trx,
+					);
+					if (savedMessage?.status === 'cancelled') {
+						return;
+					}
+
+					// Otherwise mark them as errored
+					await this.messageRepository.updateChatMessage(
+						message.id,
+						{
+							status: 'error',
+						},
+						trx,
+					);
+				});
+			},
+		});
+
+		const transform = (text: string) => {
+			const trimmed = text.trim();
+			if (!trimmed) return text;
+
+			let chunk: StructuredChunk | null = null;
+			try {
+				chunk = jsonParse<StructuredChunk>(trimmed);
+			} catch {
+				return text;
 			}
+
+			const message = aggregator.ingest(chunk);
+			const enriched: EnrichedStructuredChunk = {
+				...chunk,
+				metadata: {
+					...chunk.metadata,
+					messageId: message.id,
+					previousMessageId: message.previousMessageId,
+					retryOfMessageId: message.retryOfMessageId,
+				},
+			};
+
+			return JSON.stringify(enriched) + '\n';
 		};
 
-		const stream = captureResponseWrites(res, onChunk);
+		const stream = interceptResponseWrites(res, transform);
+
+		stream.on('finish', aggregator.finalizeAll);
+		stream.on('close', aggregator.finalizeAll);
+
 		stream.writeHead(200, JSONL_STREAM_HEADERS);
 		stream.flushHeaders();
 
-		const { executionId } = await this.workflowExecutionService.executeManually(
+		const execution = await this.workflowExecutionService.executeManually(
 			{
 				workflowData,
 				triggerToStartFrom,
@@ -897,21 +998,12 @@ export class ChatHubService {
 			true,
 			stream,
 		);
+		executionId = execution.executionId;
+
 		if (!executionId) {
 			throw new OperationalError('There was a problem starting the chat execution.');
 		}
 
-		await this.saveAIMessage({
-			id: replyId,
-			sessionId,
-			executionId,
-			previousMessageId,
-			message: partialMessage,
-			selectedModel,
-			retryOfMessageId,
-			status: 'running',
-		});
-
 		try {
 			let result: IRun | undefined;
 			try {
@@ -921,20 +1013,7 @@ export class ChatHubService {
 				}
 			} catch (error: unknown) {
 				if (error instanceof ManualExecutionCancelledError) {
-					const execution = await this.executionRepository.findWithUnflattenedData(executionId, [
-						workflowData.id,
-					]);
-					if (!execution) {
-						throw new OperationalError(`Could not find execution with ID ${executionId}`);
-					}
-
-					if (execution.status === 'canceled') {
-						await this.messageRepository.updateChatMessage(replyId, {
-							content: partialMessage || 'Generation cancelled.',
-							status: 'cancelled',
-						});
-						return;
-					}
+					return;
 				}
 
 				throw error;
@@ -952,29 +1031,52 @@ export class ChatHubService {
 				throw new OperationalError(message);
 			}
 
-			// TODO: We should consider can we just save the output from the captured stream always instead
-			// of parsing it from execution data, which seems error prone, especially with custom workflows.
-			// That could make handling multiple agents, multiple runes, tool executions etc easier...?
-			// const output = this.getAIOutput(execution, NODE_NAMES.REPLY_AGENT);
-			// if (!output) {
-			// 	throw new OperationalError('No response generated');
-			// }
-
-			await this.messageRepository.updateChatMessage(replyId, {
-				content: partialMessage,
-				status: 'success',
-			});
-
+			// TODO: Getting the title from the execution like this
+			// only works on base chat workflows. Custom agent workflows
+			// would need a different mechanism.
 			const title = this.getAIOutput(execution, NODE_NAMES.TITLE_GENERATOR_AGENT);
 			if (title) {
 				await this.sessionRepository.updateChatTitle(sessionId, title);
 			}
 		} catch (error: unknown) {
-			const message = error instanceof Error ? error.message : 'Unknown error';
-			await this.messageRepository.updateChatMessage(replyId, {
-				content: `Error: ${message}`,
-				status: 'error',
-			});
+			if (error instanceof Error) {
+				this.logger.error(`Error during chat workflow execution: ${error}`);
+			}
+			throw error;
+		}
+	}
+
+	private async executeChatWorkflowWithCleanup(
+		res: Response,
+		user: User,
+		workflow: {
+			workflowData: IWorkflowBase;
+			triggerToStartFrom: { name: string; data?: ITaskData };
+		},
+		sessionId: ChatSessionId,
+		previousMessageId: ChatMessageId,
+		selectedModel: ModelWithCredentials,
+		provider: ChatHubProvider,
+		retryOfMessageId: ChatMessageId | null = null,
+	) {
+		try {
+			await this.executeChatWorkflow(
+				res,
+				user,
+				workflow,
+				sessionId,
+				previousMessageId,
+				selectedModel,
+				retryOfMessageId,
+			);
+		} finally {
+			if (provider !== 'n8n') {
+				// TODO: If we don't wait for a bit then a followup workflow insights query
+				// that happens after executions might fail to find the workflow.
+				// Once/if we add the new workflow flag to keep these WFs around this wouldn't be needed.
+				await new Promise((resolve) => setTimeout(resolve, 3000));
+				await this.deleteChatWorkflow(workflow.workflowData.id);
+			}
 		}
 	}
 
@@ -1046,20 +1148,23 @@ export class ChatHubService {
 					mode: 'insert',
 					insertMode: 'override',
 					messages: {
-						messageValues: history.map((message) => {
-							const typeMap: Record<string, MessageRecord['type']> = {
-								human: 'user',
-								ai: 'ai',
-								system: 'system',
-							};
+						messageValues: history
+							// Empty messages can't be restored by the memory manager
+							.filter((message) => message.content.length > 0)
+							.map((message) => {
+								const typeMap: Record<string, MessageRecord['type']> = {
+									human: 'user',
+									ai: 'ai',
+									system: 'system',
+								};
 
-							// TODO: Tool messages ?
-							return {
-								type: typeMap[message.type] || 'system',
-								message: message.content,
-								hideFromUI: false,
-							};
-						}),
+								// TODO: Tool messages etc?
+								return {
+									type: typeMap[message.type] || 'system',
+									message: message.content,
+									hideFromUI: false,
+								};
+							}),
 					},
 				},
 				type: '@n8n/n8n-nodes-langchain.memoryManager',
@@ -1205,18 +1310,18 @@ export class ChatHubService {
 		sessionId,
 		executionId,
 		previousMessageId,
-		message,
+		content,
 		selectedModel,
 		retryOfMessageId,
 		status,
 	}: {
 		id: ChatMessageId;
 		sessionId: ChatSessionId;
-		executionId: string;
-		previousMessageId: ChatMessageId;
-		message: string;
+		previousMessageId: ChatMessageId | null;
+		content: string;
 		selectedModel: ModelWithCredentials;
-		retryOfMessageId?: ChatMessageId;
+		executionId?: string;
+		retryOfMessageId: ChatMessageId | null;
 		editOfMessageId?: ChatMessageId;
 		status?: ChatHubMessageStatus;
 	}) {
@@ -1224,11 +1329,11 @@ export class ChatHubService {
 			id,
 			sessionId,
 			previousMessageId,
-			executionId: parseInt(executionId, 10),
+			executionId: executionId ? parseInt(executionId, 10) : null,
 			type: 'ai',
 			name: 'AI',
 			status,
-			content: message,
+			content,
 			retryOfMessageId,
 			...selectedModel,
 		});
@@ -1470,11 +1575,11 @@ export class ChatHubService {
 		sessionId: ChatSessionId,
 		updates: {
 			title?: string;
-			credentialId?: string;
+			credentialId?: string | null;
 			provider?: ChatHubProvider;
-			model?: string;
-			workflowId?: string;
-			agentId?: string;
+			model?: string | null;
+			workflowId?: string | null;
+			agentId?: string | null;
 			agentName?: string;
 		},
 	) {
@@ -1484,6 +1589,16 @@ export class ChatHubService {
 			throw new NotFoundError('Session not found');
 		}
 
+		if (updates.provider === 'n8n') {
+			// n8n provider only stores workflowId
+			updates.model = null;
+			updates.credentialId = null;
+			updates.agentId = null;
+		} else if (updates.provider) {
+			// Other providers can't have workflowId
+			updates.workflowId = null;
+		}
+
 		return await this.sessionRepository.updateChatSession(sessionId, updates);
 	}
 

packages/cli/src/modules/chat-hub/chat-hub.types.ts

@@ -19,7 +19,6 @@ export interface ModelWithCredentials {
 export interface BaseMessagePayload {
 	userId: string;
 	sessionId: ChatSessionId;
-	replyId: ChatMessageId;
 	model: ChatHubConversationModel;
 	credentials: INodeCredentials;
 }

packages/cli/src/modules/chat-hub/chat-message.repository.ts

@@ -16,14 +16,19 @@ export class ChatHubMessageRepository extends Repository<ChatHubMessage> {
 	}
 
 	async createChatMessage(message: Partial<ChatHubMessage>, trx?: EntityManager) {
-		return await withTransaction(this.manager, trx, async (em) => {
-			await em.insert(ChatHubMessage, message);
-			const saved = await em.findOneOrFail(ChatHubMessage, {
-				where: { id: message.id },
-			});
-			await this.chatSessionRepository.updateLastMessageAt(saved.sessionId, saved.createdAt, em);
-			return saved;
-		});
+		return await withTransaction(
+			this.manager,
+			trx,
+			async (em) => {
+				await em.insert(ChatHubMessage, message);
+				const saved = await em.findOneOrFail(ChatHubMessage, {
+					where: { id: message.id },
+				});
+				await this.chatSessionRepository.updateLastMessageAt(saved.sessionId, saved.createdAt, em);
+				return saved;
+			},
+			false,
+		);
 	}
 
 	async updateChatMessage(
@@ -31,15 +36,25 @@ export class ChatHubMessageRepository extends Repository<ChatHubMessage> {
 		fields: { status?: ChatHubMessageStatus; content?: string },
 		trx?: EntityManager,
 	) {
-		return await withTransaction(this.manager, trx, async (em) => {
-			return await em.update(ChatHubMessage, { id }, fields);
-		});
+		return await withTransaction(
+			this.manager,
+			trx,
+			async (em) => {
+				return await em.update(ChatHubMessage, { id }, fields);
+			},
+			false,
+		);
 	}
 
 	async deleteChatMessage(id: ChatMessageId, trx?: EntityManager) {
-		return await withTransaction(this.manager, trx, async (em) => {
-			return await em.delete(ChatHubMessage, { id });
-		});
+		return await withTransaction(
+			this.manager,
+			trx,
+			async (em) => {
+				return await em.delete(ChatHubMessage, { id });
+			},
+			false,
+		);
 	}
 
 	async getManyBySessionId(sessionId: string, trx?: EntityManager) {

packages/cli/src/modules/chat-hub/stream-capturer.ts

@@ -1,21 +1,24 @@
+import type { ChatMessageId } from '@n8n/api-types';
 import type { Response } from 'express';
 import type { ServerResponse } from 'http';
+import type { StructuredChunk } from 'n8n-workflow';
+import { v4 as uuidv4 } from 'uuid';
+
+import type { ChatHubMessage } from './chat-hub-message.entity';
 
 type Write = ServerResponse['write'];
 
-export type ChunkListenerCb = (chunk: string) => void;
+export type ChunkTransformer = (chunk: string) => void;
 
-export function captureResponseWrites<T extends Response>(res: T, onChunk: ChunkListenerCb): T {
+export function interceptResponseWrites<T extends Response>(
+	res: T,
+	transform: ChunkTransformer,
+): T {
 	const originalWrite = res.write.bind(res) as Write;
+	const defaultEncoding = 'utf8';
 
-	const writeListener = (chunk: string | Buffer, enc?: BufferEncoding) => {
-		try {
-			const text = Buffer.isBuffer(chunk) ? chunk.toString(enc ?? 'utf8') : String(chunk);
-			void onChunk(text);
-		} catch {
-			// Don't break the stream on listener errors
-		}
-	};
+	const toText = (data: string | Buffer, enc?: BufferEncoding) =>
+		Buffer.isBuffer(data) ? data.toString(enc ?? defaultEncoding) : String(data);
 
 	function write(chunk: string | Buffer, callbackFn?: (e?: Error | null) => void): boolean;
 	function write(
@@ -25,21 +28,22 @@ export function captureResponseWrites<T extends Response>(res: T, onChunk: Chunk
 	): boolean;
 	function write(
 		chunk: string | Buffer,
-		encodingOrCallbackFn?: BufferEncoding | ((e?: Error | null) => void),
+		encodingOrCallback?: BufferEncoding | ((e?: Error | null) => void),
 		callbackFn?: (e?: Error | null) => void,
 	): boolean {
-		// TODO: We could also change the output that gets streamed from execution engine here,
-		// perhaps injecting the messageId or other metadata into the chunks. That could make
-		// AI responding with multiple messages (tools, multiple agents etc) easier to handle?
-		if (!encodingOrCallbackFn) {
-			writeListener(chunk);
-			return originalWrite(chunk);
-		} else if (typeof encodingOrCallbackFn === 'function') {
-			writeListener(chunk);
-			return originalWrite(chunk, encodingOrCallbackFn);
+		const inputText = toText(
+			chunk,
+			typeof encodingOrCallback === 'string' ? encodingOrCallback : undefined,
+		);
+
+		const outputText = transform(inputText);
+
+		if (!encodingOrCallback) {
+			return originalWrite(outputText);
+		} else if (typeof encodingOrCallback === 'function') {
+			return originalWrite(outputText, encodingOrCallback);
 		} else {
-			writeListener(chunk, encodingOrCallbackFn);
-			return originalWrite(chunk, encodingOrCallbackFn, callbackFn);
+			return originalWrite(outputText, encodingOrCallback, callbackFn);
 		}
 	}
 
@@ -47,3 +51,112 @@ export function captureResponseWrites<T extends Response>(res: T, onChunk: Chunk
 
 	return res;
 }
+
+type MessageKey = string;
+const keyOf = (m: StructuredChunk['metadata']): MessageKey =>
+	`${m.nodeId}|${m.runIndex}|${m.itemIndex}`;
+
+export type AggregatedMessage = Pick<
+	ChatHubMessage,
+	'id' | 'previousMessageId' | 'retryOfMessageId' | 'content' | 'createdAt' | 'updatedAt' | 'status'
+>;
+
+type Handlers = {
+	onBegin?: (message: AggregatedMessage) => void;
+	onItem?: (message: AggregatedMessage, delta: string) => void;
+	onEnd?: (message: AggregatedMessage) => void;
+	onError?: (message: AggregatedMessage, errText?: string) => void;
+};
+
+export function createStructuredChunkAggregator(
+	initialPreviousMessageId: ChatMessageId,
+	retryOfMessageId: ChatMessageId | null,
+	handlers: Handlers = {},
+) {
+	const { onBegin, onItem, onEnd, onError } = handlers;
+
+	const active = new Map<MessageKey, AggregatedMessage>();
+	const activeByKey = new Map<MessageKey, AggregatedMessage>();
+
+	let previousMessageId: ChatMessageId | null = initialPreviousMessageId;
+
+	const startNew = (): AggregatedMessage => {
+		const message: AggregatedMessage = {
+			id: uuidv4(),
+			previousMessageId,
+			retryOfMessageId:
+				retryOfMessageId && previousMessageId === initialPreviousMessageId
+					? retryOfMessageId
+					: null,
+			content: '',
+			createdAt: new Date(),
+			updatedAt: new Date(),
+			status: 'running',
+		};
+		previousMessageId = message.id;
+		onBegin?.(message);
+		return message;
+	};
+
+	const ensureMessage = (key: MessageKey): AggregatedMessage => {
+		let message = activeByKey.get(key);
+		if (!message) {
+			message = startNew();
+			activeByKey.set(key, message);
+		}
+		return message;
+	};
+
+	const ingest = (chunk: StructuredChunk): AggregatedMessage => {
+		const { type, content, metadata } = chunk;
+		const key = keyOf(metadata);
+
+		if (type === 'begin') {
+			if (activeByKey.has(key)) {
+				throw new Error(`Duplicate begin for key ${key}`);
+			}
+			const message = startNew();
+			activeByKey.set(key, message);
+			return message;
+		}
+
+		if (type === 'item') {
+			const message = ensureMessage(key);
+			if (typeof content === 'string' && content.length) {
+				message.content += content;
+				onItem?.(message, content);
+			}
+			return message;
+		}
+
+		if (type === 'end') {
+			const message = ensureMessage(key);
+			message.status = 'success';
+			message.updatedAt = new Date();
+			activeByKey.delete(key);
+			onEnd?.(message);
+			return message;
+		}
+
+		const message = ensureMessage(key);
+		message.status = 'error';
+		message.updatedAt = new Date();
+		if (typeof content === 'string') {
+			message.content = (message.content ? message.content + '\n\n' : '') + content;
+		}
+
+		activeByKey.delete(key);
+		onError?.(message, content);
+		return message;
+	};
+
+	const finalizeAll = () => {
+		for (const message of active.values()) {
+			message.status = 'cancelled';
+			message.updatedAt = new Date();
+		}
+		active.clear();
+	};
+
+	return { ingest, finalizeAll };
+}

packages/frontend/editor-ui/src/features/ai/chatHub/ChatView.vue

@@ -52,6 +52,7 @@ const inputRef = useTemplateRef('inputRef');
 const sessionId = computed<string>(() =>
 	typeof route.params.id === 'string' ? route.params.id : uuidv4(),
 );
+const isResponding = computed(() => chatStore.isResponding(sessionId.value));
 const isNewSession = computed(() => sessionId.value !== route.params.id);
 const scrollableRef = useTemplateRef('scrollable');
 const scrollContainerRef = computed(() => scrollableRef.value?.parentElement ?? null);
@@ -83,61 +84,56 @@ const defaultModel = useLocalStorage<ChatHubConversationModel | null>(
 	},
 );
 
-function getUpdatedModelIfAgent(model: ChatHubConversationModel | null) {
-	if (!model || !(model.provider === 'custom-agent' && 'agentId' in model)) {
-		return model;
-	}
-
-	const agent = chatStore.getAgent(model.agentId);
-	if (!agent) {
-		// agent got deleted
-		return null;
-	}
-
-	return {
-		...model,
-		// use updated name if user edited agent name
-		name: agent.name,
-	};
-}
-
 const selectedModel = computed<ChatHubConversationModel | null>(() => {
-	if (
-		currentConversation.value?.provider &&
-		currentConversation.value?.provider !== 'n8n' &&
-		currentConversation.value?.provider !== 'custom-agent' &&
-		currentConversation.value?.model
-	) {
-		// todo check if model is still available
-		return {
-			provider: currentConversation.value.provider,
-			model: currentConversation.value.model,
-			name: currentConversation.value.model,
-		};
+	let model: ChatHubConversationModel | null = null;
+	if (currentConversation.value?.provider) {
+		if (currentConversation.value.provider === 'n8n') {
+			const n8nModel = chatStore.models?.n8n.models.find(
+				(m) =>
+					currentConversation.value &&
+					m.provider === 'n8n' &&
+					m.workflowId === currentConversation.value.workflowId,
+			);
+
+			if (!n8nModel) {
+				return null;
+			}
+
+			model = n8nModel;
+		} else if (currentConversation.value.provider === 'custom-agent') {
+			const agentModel = chatStore.models?.['custom-agent'].models.find(
+				(m) =>
+					currentConversation.value &&
+					m.provider === 'custom-agent' &&
+					m.agentId === currentConversation.value.agentId,
+			);
+
+			if (!agentModel) {
+				return null;
+			}
+
+			model = agentModel;
+		} else {
+			const chatModel = chatStore.models?.[currentConversation.value.provider].models.find(
+				(m) =>
+					currentConversation.value &&
+					currentConversation.value?.provider !== 'n8n' &&
+					currentConversation.value?.provider !== 'custom-agent' &&
+					m.provider === currentConversation.value?.provider &&
+					m.model === currentConversation.value.model,
+			);
+
+			if (!chatModel) {
+				return null;
+			}
+
+			model = chatModel;
+		}
+	} else {
+		model = defaultModel.value;
 	}
 
-	if (
-		currentConversation.value?.provider === 'custom-agent' &&
-		currentConversation?.value.agentId &&
-		currentConversation?.value.agentName
-	) {
-		return getUpdatedModelIfAgent({
-			provider: currentConversation.value.provider,
-			agentId: currentConversation.value.agentId,
-			name: currentConversation.value.agentName,
-		});
-	}
-
-	if (currentConversation.value?.provider === 'n8n' && currentConversation?.value.workflowId) {
-		// todo check if workflow is still available or updated
-		return {
-			provider: currentConversation.value.provider,
-			workflowId: currentConversation.value.workflowId,
-			name: currentConversation.value.agentName ?? 'n8n workflow', // todo remove default once this gets set
-		};
-	}
-
-	return getUpdatedModelIfAgent(defaultModel.value);
+	return model;
 });
 
 const selectedCredentials = useLocalStorage<CredentialsMap>(
@@ -237,14 +233,15 @@ watch(
 			return;
 		}
 
-		if (lastMessageId !== chatStore.streamingMessageId) {
-			scrollToBottom(chatStore.streamingMessageId !== undefined);
+		const currentMessage = chatStore.lastMessage(sessionId.value);
+		if (lastMessageId !== currentMessage?.id) {
+			scrollToBottom(currentMessage !== null);
 			return;
 		}
 
 		const message = chatStore
 			.getActiveMessages(sessionId.value)
-			.find((message) => message.id === lastMessageId);
+			.find((m) => m.id === lastMessageId);
 
 		if (message?.previousMessageId) {
 			// Scroll to user's prompt when the message is being generated
@@ -319,7 +316,7 @@ onMounted(async () => {
 function onSubmit(message: string) {
 	if (
 		!message.trim() ||
-		chatStore.isResponding ||
+		isResponding.value ||
 		!selectedModel.value ||
 		isMissingSelectedCredential.value
 	) {
@@ -366,58 +363,67 @@ function handleCancelEditMessage() {
 
 function handleEditMessage(message: ChatHubMessageDto) {
 	if (
-		chatStore.isResponding ||
+		chatStore.isResponding(message.sessionId) ||
 		!['human', 'ai'].includes(message.type) ||
 		!selectedModel.value ||
-		!credentialsId.value
+		isMissingSelectedCredential.value
 	) {
 		return;
 	}
 
 	const messageToEdit = message.revisionOfMessageId ?? message.id;
 
+	const credentials = {};
+	if (
+		selectedModel.value.provider !== 'n8n' &&
+		selectedModel.value.provider !== 'custom-agent' &&
+		credentialsId.value
+	) {
+		Object.assign(credentials, {
+			[PROVIDER_CREDENTIAL_TYPE_MAP[selectedModel.value.provider]]: {
+				id: credentialsId.value,
+				name: '',
+			},
+		});
+	}
+
 	chatStore.editMessage(
 		sessionId.value,
 		messageToEdit,
 		message.content,
 		selectedModel.value,
-		selectedModel.value.provider === 'n8n' || selectedModel.value.provider === 'custom-agent'
-			? {}
-			: {
-					[PROVIDER_CREDENTIAL_TYPE_MAP[selectedModel.value.provider]]: {
-						id: credentialsId.value,
-						name: '',
-					},
-				},
+		credentials,
 	);
 	editingMessageId.value = undefined;
 }
 
 function handleRegenerateMessage(message: ChatHubMessageDto) {
 	if (
-		chatStore.isResponding ||
+		chatStore.isResponding(message.sessionId) ||
 		message.type !== 'ai' ||
 		!selectedModel.value ||
-		!credentialsId.value
+		isMissingSelectedCredential.value
 	) {
 		return;
 	}
 
 	const messageToRetry = message.retryOfMessageId ?? message.id;
 
-	chatStore.regenerateMessage(
-		sessionId.value,
-		messageToRetry,
-		selectedModel.value,
-		selectedModel.value.provider === 'n8n' || selectedModel.value.provider === 'custom-agent'
-			? {}
-			: {
-					[PROVIDER_CREDENTIAL_TYPE_MAP[selectedModel.value.provider]]: {
-						id: credentialsId.value,
-						name: '',
-					},
-				},
-	);
+	const credentials = {};
+	if (
+		selectedModel.value.provider !== 'n8n' &&
+		selectedModel.value.provider !== 'custom-agent' &&
+		credentialsId.value
+	) {
+		Object.assign(credentials, {
+			[PROVIDER_CREDENTIAL_TYPE_MAP[selectedModel.value.provider]]: {
+				id: credentialsId.value,
+				name: '',
+			},
+		});
+	}
+
+	chatStore.regenerateMessage(sessionId.value, messageToRetry, selectedModel.value, credentials);
 }
 
 async function handleSelectModel(selection: ChatHubConversationModel) {
@@ -513,7 +519,7 @@ function closeAgentEditor() {
 						:message="message"
 						:compact="isMobileDevice"
 						:is-editing="editingMessageId === message.id"
-						:is-streaming="chatStore.streamingMessageId === message.id"
+						:is-streaming="message.status === 'running'"
 						:min-height="
 							didSubmitInCurrentSession &&
 							message.type === 'ai' &&
@@ -543,7 +549,7 @@ function closeAgentEditor() {
 					<ChatPrompt
 						ref="inputRef"
 						:class="$style.prompt"
-						:is-responding="chatStore.isResponding"
+						:is-responding="isResponding"
 						:selected-model="selectedModel"
 						:is-missing-credentials="isMissingSelectedCredential"
 						@submit="onSubmit"

packages/frontend/editor-ui/src/features/ai/chatHub/chat.api.ts

@@ -14,8 +14,8 @@ import type {
 	ChatHubCreateAgentRequest,
 	ChatHubUpdateAgentRequest,
 	ChatHubUpdateConversationRequest,
+	EnrichedStructuredChunk,
 } from '@n8n/api-types';
-import type { StructuredChunk } from 'n8n-workflow';
 
 // Workflows stream data as newline separated JSON objects (jsonl)
 const STREAM_SEPARATOR = '\n';
@@ -31,11 +31,11 @@ export const fetchChatModelsApi = async (
 export function sendMessageApi(
 	ctx: IRestApiContext,
 	payload: ChatHubSendMessageRequest,
-	onMessageUpdated: (data: StructuredChunk) => void,
+	onMessageUpdated: (data: EnrichedStructuredChunk) => void,
 	onDone: () => void,
 	onError: (e: Error) => void,
 ) {
-	void streamRequest<StructuredChunk>(
+	void streamRequest<EnrichedStructuredChunk>(
 		ctx,
 		'/chat/conversations/send',
 		payload,
@@ -51,11 +51,11 @@ export function editMessageApi(
 	sessionId: ChatSessionId,
 	editId: ChatMessageId,
 	payload: ChatHubEditMessageRequest,
-	onMessageUpdated: (data: StructuredChunk) => void,
+	onMessageUpdated: (data: EnrichedStructuredChunk) => void,
 	onDone: () => void,
 	onError: (e: Error) => void,
 ) {
-	void streamRequest<StructuredChunk>(
+	void streamRequest<EnrichedStructuredChunk>(
 		ctx,
 		`/chat/conversations/${sessionId}/messages/${editId}/edit`,
 		payload,
@@ -71,11 +71,11 @@ export function regenerateMessageApi(
 	sessionId: ChatSessionId,
 	retryId: ChatMessageId,
 	payload: ChatHubRegenerateMessageRequest,
-	onMessageUpdated: (data: StructuredChunk) => void,
+	onMessageUpdated: (data: EnrichedStructuredChunk) => void,
 	onDone: () => void,
 	onError: (e: Error) => void,
 ) {
-	void streamRequest<StructuredChunk>(
+	void streamRequest<EnrichedStructuredChunk>(
 		ctx,
 		`/chat/conversations/${sessionId}/messages/${retryId}/regenerate`,
 		payload,

packages/frontend/editor-ui/src/features/ai/chatHub/chat.store.ts

@@ -1,6 +1,6 @@
 import { defineStore } from 'pinia';
 import { CHAT_STORE } from './constants';
-import { computed, ref } from 'vue';
+import { ref } from 'vue';
 import { v4 as uuidv4 } from 'uuid';
 import {
 	fetchChatModelsApi,
@@ -32,22 +32,20 @@ import type {
 	ChatHubAgentDto,
 	ChatHubCreateAgentRequest,
 	ChatHubUpdateAgentRequest,
+	EnrichedStructuredChunk,
+	ChatHubMessageStatus,
 } from '@n8n/api-types';
 import type { CredentialsMap, ChatMessage, ChatConversation } from './chat.types';
-import type { StructuredChunk } from 'n8n-workflow';
 import { retry } from '@n8n/utils/retry';
 
 export const useChatStore = defineStore(CHAT_STORE, () => {
 	const rootStore = useRootStore();
 	const models = ref<ChatModelsResponse>();
 	const loadingModels = ref(false);
-	const streamingMessageId = ref<string>();
 	const sessions = ref<ChatHubSessionDto[]>([]);
 	const agents = ref<ChatHubAgentDto[]>([]);
 	const currentEditingAgent = ref<ChatHubAgentDto | null>(null);
 
-	const isResponding = computed(() => streamingMessageId.value !== undefined);
-
 	const conversationsBySession = ref<Map<ChatSessionId, ChatConversation>>(new Map());
 
 	const getConversation = (sessionId: ChatSessionId): ChatConversation | undefined =>
@@ -76,6 +74,21 @@ export const useChatStore = defineStore(CHAT_STORE, () => {
 		return conversation;
 	}
 
+	function lastMessage(sessionId: ChatSessionId): ChatMessage | null {
+		const conversation = getConversation(sessionId);
+		if (!conversation || conversation.activeMessageChain.length === 0) {
+			return null;
+		}
+
+		const messageId = conversation.activeMessageChain[conversation.activeMessageChain.length - 1];
+		return conversation.messages[messageId] ?? null;
+	}
+
+	function isResponding(sessionId: ChatSessionId) {
+		const last = lastMessage(sessionId);
+		return last?.status === 'running';
+	}
+
 	function computeActiveChain(
 		messages: Record<ChatMessageId, ChatMessage>,
 		messageId: ChatMessageId | null,
@@ -199,6 +212,8 @@ export const useChatStore = defineStore(CHAT_STORE, () => {
 		// TODO: Recomputing the entire graph shouldn't be needed here
 		conversation.messages = linkMessages(Object.values(conversation.messages));
 		conversation.activeMessageChain = computeActiveChain(conversation.messages, message.id);
+
+		return message;
 	}
 
 	function replaceMessageContent(
@@ -225,6 +240,21 @@ export const useChatStore = defineStore(CHAT_STORE, () => {
 		message.content += chunk;
 	}
 
+	function updateMessage(
+		sessionId: ChatSessionId,
+		messageId: ChatMessageId,
+		status: ChatHubMessageStatus,
+	) {
+		const conversation = ensureConversation(sessionId);
+		const message = conversation.messages[messageId];
+		if (!message) {
+			throw new Error(`Message with ID ${messageId} not found in session ${sessionId}`);
+		}
+
+		message.status = status;
+		message.updatedAt = new Date().toISOString();
+	}
+
 	async function fetchChatModels(credentialMap: CredentialsMap) {
 		loadingModels.value = true;
 		models.value = await fetchChatModelsApi(rootStore.restApiContext, {
@@ -255,15 +285,12 @@ export const useChatStore = defineStore(CHAT_STORE, () => {
 	}
 
 	function onBeginMessage(
-		sessionId: string,
-		messageId: string,
-		replyToMessageId: string,
-		retryOfMessageId: string | null,
-		_nodeId: string,
-		_runIndex?: number,
+		sessionId: ChatSessionId,
+		messageId: ChatMessageId,
+		previousMessageId: ChatMessageId | null,
+		retryOfMessageId: ChatMessageId | null,
+		status: ChatHubMessageStatus = 'running',
 	) {
-		streamingMessageId.value = messageId;
-
 		addMessage(sessionId, {
 			id: messageId,
 			sessionId,
@@ -275,10 +302,10 @@ export const useChatStore = defineStore(CHAT_STORE, () => {
 			workflowId: null,
 			executionId: null,
 			agentId: null,
-			status: 'success',
+			status,
 			createdAt: new Date().toISOString(),
 			updatedAt: new Date().toISOString(),
-			previousMessageId: replyToMessageId,
+			previousMessageId,
 			retryOfMessageId,
 			revisionOfMessageId: null,
 			responses: [],
@@ -286,58 +313,77 @@ export const useChatStore = defineStore(CHAT_STORE, () => {
 		});
 	}
 
-	function onChunk(
-		sessionId: string,
-		messageId: string,
-		chunk: string,
-		_nodeId?: string,
-		_runIndex?: number,
-	) {
+	function ensureMessage(
+		sessionId: ChatSessionId,
+		messageId: ChatMessageId,
+		previousMessageId: ChatMessageId | null,
+		retryOfMessageId: ChatMessageId | null,
+	): ChatMessage {
+		const conversation = ensureConversation(sessionId);
+		const message = conversation.messages[messageId];
+		if (message) {
+			return message;
+		}
+
+		return addMessage(sessionId, {
+			id: messageId,
+			sessionId,
+			type: 'ai',
+			name: 'AI',
+			content: '',
+			provider: null,
+			model: null,
+			workflowId: null,
+			executionId: null,
+			status: 'running',
+			createdAt: new Date().toISOString(),
+			updatedAt: new Date().toISOString(),
+			previousMessageId,
+			retryOfMessageId,
+			revisionOfMessageId: null,
+			responses: [],
+			alternatives: [],
+			agentId: null,
+		});
+	}
+
+	function onChunk(sessionId: string, messageId: string, chunk: string) {
 		appendMessage(sessionId, messageId, chunk);
 	}
 
-	function onEndMessage() {
-		streamingMessageId.value = undefined;
+	function onEndMessage(sessionId: ChatSessionId, messageId: ChatMessageId) {
+		updateMessage(sessionId, messageId, 'success');
 	}
 
-	function onStreamMessage(
-		sessionId: string,
-		message: StructuredChunk,
-		messageId: string,
-		replyToMessageId: string,
-		retryOfMessageId: string | null,
-	) {
-		const nodeId = message.metadata?.nodeId || 'unknown';
-		const runIndex = message.metadata?.runIndex;
+	function onStreamMessage(sessionId: string, chunk: EnrichedStructuredChunk) {
+		const { messageId, previousMessageId, retryOfMessageId } = chunk.metadata;
 
-		switch (message.type) {
+		switch (chunk.type) {
 			case 'begin':
-				onBeginMessage(sessionId, messageId, replyToMessageId, retryOfMessageId, nodeId, runIndex);
+				onBeginMessage(sessionId, messageId, previousMessageId, retryOfMessageId);
 				break;
 			case 'item':
-				onChunk(sessionId, messageId, message.content ?? '', nodeId, runIndex);
+				onChunk(sessionId, messageId, chunk.content ?? '');
 				break;
 			case 'end':
-				onEndMessage();
+				onEndMessage(sessionId, messageId);
 				break;
-			case 'error':
-				if (streamingMessageId.value === messageId) {
-					onChunk(
-						sessionId,
-						messageId,
-						`Error: ${message.content ?? 'Unknown error'}`,
-						nodeId,
-						runIndex,
-					);
-					onEndMessage();
+			case 'error': {
+				// Ignore errors after cancellation
+				const message = ensureMessage(sessionId, messageId, previousMessageId, retryOfMessageId);
+
+				if (message.status === 'cancelled') {
+					return;
 				}
+
+				updateMessage(sessionId, messageId, 'error');
+				onChunk(sessionId, messageId, message.content ?? '');
 				break;
+			}
 		}
 	}
 
-	async function onStreamDone(sessionId: string) {
-		streamingMessageId.value = undefined;
-
+	async function onStreamDone(sessionId: ChatSessionId) {
 		// wait up to 3 seconds until conversation title is generated
 		await retry(
 			async () => {
@@ -354,8 +400,19 @@ export const useChatStore = defineStore(CHAT_STORE, () => {
 	}
 
 	// eslint-disable-next-line @typescript-eslint/no-unused-vars
-	function onStreamError(_e: Error) {
-		streamingMessageId.value = undefined;
+	function onStreamError(sessionId: ChatSessionId, _e: Error) {
+		const conversation = getConversation(sessionId);
+		if (!conversation) {
+			return;
+		}
+
+		// TODO: Not sure if we want to mark all running messages as errored?
+		for (const messageId of conversation.activeMessageChain) {
+			const message = conversation.messages[messageId];
+			if (message.status === 'running') {
+				updateMessage(sessionId, messageId, 'error');
+			}
+		}
 	}
 
 	function sendMessage(
@@ -365,7 +422,6 @@ export const useChatStore = defineStore(CHAT_STORE, () => {
 		credentials: ChatHubSendMessageRequest['credentials'],
 	) {
 		const messageId = uuidv4();
-		const replyId = uuidv4();
 		const conversation = ensureConversation(sessionId);
 		const previousMessageId = conversation.activeMessageChain.length
 			? conversation.activeMessageChain[conversation.activeMessageChain.length - 1]
@@ -398,14 +454,13 @@ export const useChatStore = defineStore(CHAT_STORE, () => {
 				model,
 				messageId,
 				sessionId,
-				replyId,
 				message,
 				credentials,
 				previousMessageId,
 			},
-			(chunk: StructuredChunk) => onStreamMessage(sessionId, chunk, replyId, messageId, null),
+			(chunk: EnrichedStructuredChunk) => onStreamMessage(sessionId, chunk),
 			async () => await onStreamDone(sessionId),
-			onStreamError,
+			(e) => onStreamError(sessionId, e),
 		);
 	}
 
@@ -417,7 +472,6 @@ export const useChatStore = defineStore(CHAT_STORE, () => {
 		credentials: ChatHubSendMessageRequest['credentials'],
 	) {
 		const messageId = uuidv4();
-		const replyId = uuidv4();
 
 		const conversation = ensureConversation(sessionId);
 		const message = conversation.messages[editId];
@@ -455,13 +509,12 @@ export const useChatStore = defineStore(CHAT_STORE, () => {
 			{
 				model,
 				messageId,
-				replyId,
 				message: content,
 				credentials,
 			},
-			(chunk: StructuredChunk) => onStreamMessage(sessionId, chunk, replyId, messageId, null),
+			(chunk: EnrichedStructuredChunk) => onStreamMessage(sessionId, chunk),
 			async () => await onStreamDone(sessionId),
-			onStreamError,
+			(e) => onStreamError(sessionId, e),
 		);
 	}
 
@@ -471,7 +524,6 @@ export const useChatStore = defineStore(CHAT_STORE, () => {
 		model: ChatHubConversationModel,
 		credentials: ChatHubSendMessageRequest['credentials'],
 	) {
-		const replyId = uuidv4();
 		const conversation = ensureConversation(sessionId);
 		const previousMessageId = conversation.messages[retryId]?.previousMessageId ?? null;
 
@@ -485,26 +537,24 @@ export const useChatStore = defineStore(CHAT_STORE, () => {
 			retryId,
 			{
 				model,
-				replyId,
 				credentials,
 			},
-			(chunk: StructuredChunk) =>
-				onStreamMessage(sessionId, chunk, replyId, previousMessageId, retryId),
+			(chunk: EnrichedStructuredChunk) => onStreamMessage(sessionId, chunk),
 			async () => await onStreamDone(sessionId),
-			onStreamError,
+			(e) => onStreamError(sessionId, e),
 		);
 	}
 
 	async function stopStreamingMessage(sessionId: ChatSessionId) {
-		if (streamingMessageId.value) {
-			const messageId = streamingMessageId.value;
-			onEndMessage();
-			await stopGenerationApi(rootStore.restApiContext, sessionId, messageId);
+		const currentMessage = lastMessage(sessionId);
+
+		if (currentMessage && currentMessage.status === 'running') {
+			updateMessage(sessionId, currentMessage.id, 'cancelled');
+			await stopGenerationApi(rootStore.restApiContext, sessionId, currentMessage.id);
 		}
 	}
 
 	function updateSession(sessionId: ChatSessionId, toUpdate: Partial<ChatHubSessionDto>) {
-		console.log(JSON.stringify(sessions.value.find((s) => s.id === sessionId)));
 		sessions.value = sessions.value.map((session) =>
 			session.id === sessionId
 				? {
@@ -513,7 +563,6 @@ export const useChatStore = defineStore(CHAT_STORE, () => {
 					}
 				: session,
 		);
-		console.log(JSON.stringify(sessions.value.find((s) => s.id === sessionId)));
 	}
 
 	async function renameSession(sessionId: ChatSessionId, title: string) {
@@ -613,8 +662,8 @@ export const useChatStore = defineStore(CHAT_STORE, () => {
 		conversationsBySession,
 		loadingModels,
 		isResponding,
-		streamingMessageId,
 		getAgent,
+		lastMessage,
 		fetchChatModels,
 		updateSessionModel,
 		sendMessage,

packages/frontend/editor-ui/src/features/ai/chatHub/components/ChatMessage.vue

@@ -199,7 +199,11 @@ onBeforeMount(() => {
 					<VueMarkdown
 						:key="forceReRenderKey"
 						:class="[$style.chatMessageMarkdown, 'chat-message-markdown']"
-						:source="message.content"
+						:source="
+							message.status === 'error' && !message.content
+								? 'Error: Unknown error occurred'
+								: message.content
+						"
 						:options="markdownOptions"
 						:plugins="[linksNewTabPlugin]"
 					/>