Add working version

packages/editor-ui/src/components/Node/NodeCreator/composables/useActions.ts

@@ -43,7 +43,7 @@ import { useExternalHooks } from '@/composables/useExternalHooks';
 import { sortNodeCreateElements, transformNodeType } from '../utils';
 import { useI18n } from '@/composables/useI18n';
 import { useCanvasStore } from '@/stores/canvas.store';
-import { adjustNewlyConnectedNodes } from '@/utils/connectionNodeUtils';
+import { adjustNewNodes } from '@/utils/connectionNodeUtils';
 
 export const useActions = () => {
 	const nodeCreatorStore = useNodeCreatorStore();
@@ -287,7 +287,7 @@ export const useActions = () => {
 		}
 
 		if (addedNodes.length === 2) {
-			adjustNewlyConnectedNodes(addedNodes[0], addedNodes[1]);
+			adjustNewNodes(addedNodes[0], addedNodes[1]);
 		}
 
 		addedNodes.forEach((node, index) => {

packages/editor-ui/src/composables/useCanvasOperations.ts

@@ -96,7 +96,7 @@ import type { useRouter } from 'vue-router';
 import { useClipboard } from '@/composables/useClipboard';
 import { useUniqueNodeName } from '@/composables/useUniqueNodeName';
 import { isPresent } from '../utils/typesUtils';
-import { adjustNewlyConnectedNodes } from '@/utils/connectionNodeUtils';
+import { adjustNewNodes } from '@/utils/connectionNodeUtils';
 
 type AddNodeData = Partial<INodeUi> & {
 	type: string;
@@ -673,37 +673,46 @@ export function useCanvasOperations({ router }: { router: ReturnType<typeof useR
 			});
 
 			if (mode === CanvasConnectionMode.Input) {
-				createConnection({
-					source: nodeId,
-					sourceHandle: nodeHandle,
-					target: lastInteractedWithNodeId,
-					targetHandle: lastInteractedWithNodeHandle,
-				});
+				createConnection(
+					{
+						source: nodeId,
+						sourceHandle: nodeHandle,
+						target: lastInteractedWithNodeId,
+						targetHandle: lastInteractedWithNodeHandle,
+					},
+					{ parentIsNew: true },
+				);
 			} else {
-				createConnection({
-					source: lastInteractedWithNodeId,
-					sourceHandle: lastInteractedWithNodeHandle,
-					target: nodeId,
-					targetHandle: nodeHandle,
-				});
+				createConnection(
+					{
+						source: lastInteractedWithNodeId,
+						sourceHandle: lastInteractedWithNodeHandle,
+						target: nodeId,
+						targetHandle: nodeHandle,
+					},
+					{ childIsNew: true },
+				);
 			}
 		} else {
 			// If a node is last selected then connect between the active and its child ones
 			// Connect active node to the newly created one
-			createConnection({
-				source: lastInteractedWithNodeId,
-				sourceHandle: createCanvasConnectionHandleString({
-					mode: CanvasConnectionMode.Output,
-					type: NodeConnectionType.Main,
-					index: 0,
-				}),
-				target: node.id,
-				targetHandle: createCanvasConnectionHandleString({
-					mode: CanvasConnectionMode.Input,
-					type: NodeConnectionType.Main,
-					index: 0,
-				}),
-			});
+			createConnection(
+				{
+					source: lastInteractedWithNodeId,
+					sourceHandle: createCanvasConnectionHandleString({
+						mode: CanvasConnectionMode.Output,
+						type: NodeConnectionType.Main,
+						index: 0,
+					}),
+					target: node.id,
+					targetHandle: createCanvasConnectionHandleString({
+						mode: CanvasConnectionMode.Input,
+						type: NodeConnectionType.Main,
+						index: 0,
+					}),
+				},
+				{ childIsNew: true },
+			);
 		}
 
 		if (lastInteractedWithNodeConnection) {
@@ -711,16 +720,19 @@ export function useCanvasOperations({ router }: { router: ReturnType<typeof useR
 
 			const targetNode = workflowsStore.getNodeById(lastInteractedWithNodeConnection.target);
 			if (targetNode) {
-				createConnection({
-					source: node.id,
-					sourceHandle: createCanvasConnectionHandleString({
-						mode: CanvasConnectionMode.Input,
-						type: NodeConnectionType.Main,
-						index: 0,
-					}),
-					target: lastInteractedWithNodeConnection.target,
-					targetHandle: lastInteractedWithNodeConnection.targetHandle,
-				});
+				createConnection(
+					{
+						source: node.id,
+						sourceHandle: createCanvasConnectionHandleString({
+							mode: CanvasConnectionMode.Input,
+							type: NodeConnectionType.Main,
+							index: 0,
+						}),
+						target: lastInteractedWithNodeConnection.target,
+						targetHandle: lastInteractedWithNodeConnection.targetHandle,
+					},
+					{ parentIsNew: true },
+				);
 			}
 		}
 	}
@@ -1079,7 +1091,7 @@ export function useCanvasOperations({ router }: { router: ReturnType<typeof useR
 
 	function createConnection(
 		connection: Connection,
-		{ trackHistory = false, keepPristine = false } = {},
+		{ trackHistory = false, keepPristine = false, parentIsNew = false, childIsNew = false } = {},
 	) {
 		const sourceNode = workflowsStore.getNodeById(connection.source);
 		const targetNode = workflowsStore.getNodeById(connection.target);
@@ -1105,7 +1117,7 @@ export function useCanvasOperations({ router }: { router: ReturnType<typeof useR
 			return;
 		}
 
-		adjustNewlyConnectedNodes(sourceNode, targetNode);
+		adjustNewNodes(sourceNode, targetNode, { parentIsNew, childIsNew });
 
 		workflowsStore.addConnection({
 			connection: mappedConnection,

packages/editor-ui/src/utils/connectionNodeUtils.test.ts

@@ -1,39 +1,76 @@
+import { useWorkflowsStore } from '@/stores/workflows.store';
 import { AGENT_NODE_TYPE, CHAT_TRIGGER_NODE_TYPE, MANUAL_TRIGGER_NODE_TYPE } from '@/constants';
-import { adjustNewlyConnectedNodes } from './connectionNodeUtils';
+import { adjustNewNodes } from '@/utils/connectionNodeUtils';
+import { createPinia, setActivePinia } from 'pinia';
+
+const getParentNodesByDepth = vi.fn();
+const getNode = vi.fn();
+vi.mock('@/stores/workflow.store', () => ({
+	useWorkflowsStore: vi.fn(() => ({
+		getParentNodesByDepth,
+		getNode,
+	})),
+}));
 
 describe('adjustNewlyConnectedNodes', () => {
-	it('modifies promptType with ChatTrigger->Agent', () => {
+	beforeEach(() => {
+		setActivePinia(createPinia());
+	});
+
+	afterEach(() => {
+		vi.clearAllMocks();
+	});
+
+	it('modifies promptType with ChatTrigger->new Agent', () => {
 		const parent = { type: CHAT_TRIGGER_NODE_TYPE };
 		const child = { type: AGENT_NODE_TYPE };
-		adjustNewlyConnectedNodes(parent, child);
+		adjustNewNodes(parent, child, { parentIsNew: false });
 		expect(child).toEqual({
 			type: AGENT_NODE_TYPE,
 		});
 	});
-	it('does not modify promptType with ManualTrigger->Agent', () => {
+
+	it('modifies promptType with new ChatTrigger->new Agent', () => {
+		const parent = { type: CHAT_TRIGGER_NODE_TYPE };
+		const child = { type: AGENT_NODE_TYPE };
+		adjustNewNodes(parent, child);
+		expect(child).toEqual({
+			type: AGENT_NODE_TYPE,
+		});
+	});
+
+	it('does not modify promptType with ManualTrigger->new Agent', () => {
 		const parent = { type: MANUAL_TRIGGER_NODE_TYPE };
 		const child = { type: AGENT_NODE_TYPE };
-		adjustNewlyConnectedNodes(parent, child);
+		adjustNewNodes(parent, child, { parentIsNew: false });
 		expect(child).toEqual({
 			type: AGENT_NODE_TYPE,
 			parameters: { promptType: 'define' },
 		});
 	});
 
-	it('modifies sessionId with ChatTrigger->Memory', () => {
-		const parent = { type: CHAT_TRIGGER_NODE_TYPE };
-		const child = { type: '@n8n/n8n-nodes-langchain.memoryBufferWindow' };
-		adjustNewlyConnectedNodes(parent, child);
-		expect(child).toEqual({
+	it('modifies sessionId with ChatTrigger->(new Memory->Agent)', () => {
+		const trigger = { type: CHAT_TRIGGER_NODE_TYPE, name: 'trigger' };
+		getParentNodesByDepth.mockReturnValue([{ name: trigger.name }]);
+		getNode.mockReturnValue({ type: trigger.type });
+
+		const child = { type: AGENT_NODE_TYPE };
+		const parent = { type: '@n8n/n8n-nodes-langchain.memoryBufferWindow' };
+		adjustNewNodes(parent, child, { childIsNew: false });
+		expect(parent).toEqual({
 			type: '@n8n/n8n-nodes-langchain.memoryBufferWindow',
 		});
 	});
 
-	it('does not modify sessionId with ManualTrigger->Memory', () => {
-		const parent = { type: MANUAL_TRIGGER_NODE_TYPE };
-		const child = { type: '@n8n/n8n-nodes-langchain.memoryBufferWindow' };
-		adjustNewlyConnectedNodes(parent, child);
-		expect(child).toEqual({
+	it('does not modify sessionId with ManualTrigger->(new Memory->Agent)', () => {
+		const trigger = { type: MANUAL_TRIGGER_NODE_TYPE, name: 'trigger' };
+		getParentNodesByDepth.mockReturnValue([{ name: trigger.name }]);
+		getNode.mockReturnValue({ type: trigger.type });
+
+		const child = { type: AGENT_NODE_TYPE, name: 'myAgent' };
+		const parent = { type: '@n8n/n8n-nodes-langchain.memoryBufferWindow' };
+		adjustNewNodes(parent, child, { childIsNew: false });
+		expect(parent).toEqual({
 			type: '@n8n/n8n-nodes-langchain.memoryBufferWindow',
 			parameters: { sessionIdType: 'customKey' },
 		});

packages/editor-ui/src/utils/connectionNodeUtils.ts

@@ -7,8 +7,8 @@ import {
 	OPEN_AI_NODE_MESSAGE_ASSISTANT_TYPE,
 	QA_CHAIN_NODE_TYPE,
 } from '@/constants';
-import { getParentNodes } from '@/components/ButtonParameter/utils';
 import { useWorkflowsStore } from '@/stores/workflows.store';
+import type { AddedNode } from '@/Interface';
 
 const AI_NODES = [
 	QA_CHAIN_NODE_TYPE,
@@ -29,25 +29,37 @@ const MEMORY_NODE_NAMES = [
 
 const PROMPT_PROVIDER_NODE_NAMES = [CHAT_TRIGGER_NODE_TYPE];
 
-type NodeWithType = Pick<INode, 'type'>;
-
-const { getCurrentWorkflow, getNodeByName } = useWorkflowsStore();
-
-export function adjustNewlyConnectedNodes(parent: INode, child: INode) {
-	const workflow = getCurrentWorkflow();
-
-	if (workflow.getParentNodesByDepth(child.name, 1).length > 0) {
-		return;
-	}
+export function adjustNewNodes(
+	parent: AddedNode,
+	child: AddedNode,
+	{ parentIsNew = true, childIsNew = true } = {},
+) {
+	if (childIsNew) adjustNewChild(parent, child);
+	if (parentIsNew) adjustNewParent(parent, child);
+}
 
+function adjustNewChild(parent: AddedNode, child: AddedNode) {
 	if (!PROMPT_PROVIDER_NODE_NAMES.includes(parent.type) && AI_NODES.includes(child.type)) {
-		Object.assign<INode, Partial<INode>>(child, {
+		Object.assign<AddedNode, Partial<INode>>(child, {
 			parameters: { promptType: 'define' },
 		});
 	}
-	if (!PROMPT_PROVIDER_NODE_NAMES.includes(parent.type) && MEMORY_NODE_NAMES.includes(child.type)) {
-		Object.assign<INode, Partial<INode>>(child, {
-			parameters: { sessionIdType: 'customKey' },
-		});
+}
+
+function adjustNewParent(parent: AddedNode, child: AddedNode) {
+	if (MEMORY_NODE_NAMES.includes(parent.type) && child.name) {
+		const { getCurrentWorkflow } = useWorkflowsStore();
+		const workflow = getCurrentWorkflow();
+
+		// If a memory node is added to an Agent, the memory node is actually the parent since it provides input
+		// So we need to look for the Agent's parents to determine if it's a prompt provider
+		const ps = workflow.getParentNodesByDepth(child.name, 1);
+		if (
+			!ps.some((x) => PROMPT_PROVIDER_NODE_NAMES.includes(workflow.getNode(x.name)?.type ?? ''))
+		) {
+			Object.assign<AddedNode, Partial<INode>>(parent, {
+				parameters: { sessionIdType: 'customKey' },
+			});
+		}
 	}
 }