From dd3b2cb62da08980ddf5a385b282f8ce189767e5 Mon Sep 17 00:00:00 2001
From: Alex Grozav <alex@grozav.com>
Date: Wed, 25 Sep 2024 12:59:54 +0300
Subject: [PATCH] fix(editor): Fix node deletion undo/redo in new canvas
 (no-changelog) (#10935)

---
 .../__tests__/useCanvasOperations.spec.ts     | 59 +++++++++++-
 .../src/composables/useCanvasOperations.ts    | 95 ++++++++++++++++++-
 .../editor-ui/src/stores/workflows.store.ts   |  2 -
 3 files changed, 148 insertions(+), 8 deletions(-)

diff --git a/packages/editor-ui/src/composables/__tests__/useCanvasOperations.spec.ts b/packages/editor-ui/src/composables/__tests__/useCanvasOperations.spec.ts
index 0a6b75d1e..4a329201a 100644
--- a/packages/editor-ui/src/composables/__tests__/useCanvasOperations.spec.ts
+++ b/packages/editor-ui/src/composables/__tests__/useCanvasOperations.spec.ts
@@ -616,7 +616,6 @@ describe('useCanvasOperations', () => {
 			deleteNode(id, { trackHistory: true });
 
 			expect(workflowsStore.removeNodeById).toHaveBeenCalledWith(id);
-			expect(workflowsStore.removeNodeConnectionsById).toHaveBeenCalledWith(id);
 			expect(workflowsStore.removeNodeExecutionDataById).toHaveBeenCalledWith(id);
 			expect(historyStore.pushCommandToUndo).toHaveBeenCalledWith(new RemoveNodeCommand(node));
 		});
@@ -644,7 +643,6 @@ describe('useCanvasOperations', () => {
 			deleteNode(id, { trackHistory: false });
 
 			expect(workflowsStore.removeNodeById).toHaveBeenCalledWith(id);
-			expect(workflowsStore.removeNodeConnectionsById).toHaveBeenCalledWith(id);
 			expect(workflowsStore.removeNodeExecutionDataById).toHaveBeenCalledWith(id);
 			expect(historyStore.pushCommandToUndo).not.toHaveBeenCalled();
 		});
@@ -714,7 +712,6 @@ describe('useCanvasOperations', () => {
 			deleteNode(nodes[1].id);
 
 			expect(workflowsStore.removeNodeById).toHaveBeenCalledWith(nodes[1].id);
-			expect(workflowsStore.removeNodeConnectionsById).toHaveBeenCalledWith(nodes[1].id);
 			expect(workflowsStore.removeNodeExecutionDataById).toHaveBeenCalledWith(nodes[1].id);
 			expect(workflowsStore.removeNodeById).toHaveBeenCalledWith(nodes[1].id);
 		});
@@ -1356,6 +1353,62 @@ describe('useCanvasOperations', () => {
 		});
 	});
 
+	describe('deleteConnectionsByNodeId', () => {
+		it('should delete all connections for a given node ID', () => {
+			const workflowsStore = mockedStore(useWorkflowsStore);
+			const { deleteConnectionsByNodeId } = useCanvasOperations({ router });
+
+			const node1 = createTestNode({ id: 'node1', name: 'Node 1' });
+			const node2 = createTestNode({ id: 'node2', name: 'Node 1' });
+
+			workflowsStore.workflow.connections = {
+				[node1.name]: {
+					[NodeConnectionType.Main]: [
+						[{ node: node2.name, type: NodeConnectionType.Main, index: 0 }],
+					],
+				},
+				node2: {
+					[NodeConnectionType.Main]: [
+						[{ node: node1.name, type: NodeConnectionType.Main, index: 0 }],
+					],
+				},
+			};
+
+			workflowsStore.getNodeById.mockReturnValue(node1);
+			workflowsStore.getNodeByName.mockReturnValueOnce(node1).mockReturnValueOnce(node2);
+
+			deleteConnectionsByNodeId(node1.id);
+
+			expect(workflowsStore.removeConnection).toHaveBeenCalledWith({
+				connection: [
+					{ node: node1.name, type: NodeConnectionType.Main, index: 0 },
+					{ node: node2.name, type: NodeConnectionType.Main, index: 0 },
+				],
+			});
+
+			expect(workflowsStore.removeConnection).toHaveBeenCalledWith({
+				connection: [
+					{ node: node2.name, type: NodeConnectionType.Main, index: 0 },
+					{ node: node1.name, type: NodeConnectionType.Main, index: 0 },
+				],
+			});
+
+			expect(workflowsStore.workflow.connections[node1.name]).toBeUndefined();
+		});
+
+		it('should not delete connections if node ID does not exist', () => {
+			const workflowsStore = mockedStore(useWorkflowsStore);
+			const { deleteConnectionsByNodeId } = useCanvasOperations({ router });
+
+			const nodeId = 'nonexistent';
+			workflowsStore.getNodeById.mockReturnValue(undefined);
+
+			deleteConnectionsByNodeId(nodeId);
+
+			expect(workflowsStore.removeConnection).not.toHaveBeenCalled();
+		});
+	});
+
 	describe('duplicateNodes', () => {
 		it('should duplicate nodes', async () => {
 			const workflowsStore = mockedStore(useWorkflowsStore);
diff --git a/packages/editor-ui/src/composables/useCanvasOperations.ts b/packages/editor-ui/src/composables/useCanvasOperations.ts
index c5ddbd4b9..ca268b645 100644
--- a/packages/editor-ui/src/composables/useCanvasOperations.ts
+++ b/packages/editor-ui/src/composables/useCanvasOperations.ts
@@ -236,7 +236,7 @@ export function useCanvasOperations({ router }: { router: ReturnType<typeof useR
 		await renameNode(currentName, previousName);
 	}
 
-	function connectAdjacentNodes(id: string) {
+	function connectAdjacentNodes(id: string, { trackHistory = false } = {}) {
 		const node = workflowsStore.getNodeById(id);
 
 		if (!node) {
@@ -262,6 +262,23 @@ export function useCanvasOperations({ router }: { router: ReturnType<typeof useR
 
 					if (!outgoingNodeId) continue;
 
+					if (trackHistory) {
+						historyStore.pushCommandToUndo(
+							new AddConnectionCommand([
+								{
+									node: incomingConnection.node,
+									type,
+									index: 0,
+								},
+								{
+									node: outgoingConnection.node,
+									type,
+									index: 0,
+								},
+							]),
+						);
+					}
+
 					createConnection({
 						source: incomingNodeId,
 						sourceHandle: createCanvasConnectionHandleString({
@@ -289,8 +306,13 @@ export function useCanvasOperations({ router }: { router: ReturnType<typeof useR
 			historyStore.startRecordingUndo();
 		}
 
-		connectAdjacentNodes(id);
-		workflowsStore.removeNodeConnectionsById(id);
+		if (uiStore.lastInteractedWithNodeId === id) {
+			uiStore.lastInteractedWithNodeId = null;
+		}
+
+		connectAdjacentNodes(id, { trackHistory });
+		deleteConnectionsByNodeId(id, { trackHistory, trackBulk: false });
+
 		workflowsStore.removeNodeExecutionDataById(id);
 		workflowsStore.removeNodeById(id);
 
@@ -1135,6 +1157,72 @@ export function useCanvasOperations({ router }: { router: ReturnType<typeof useR
 		deleteConnection(mapLegacyConnectionToCanvasConnection(sourceNode, targetNode, connection));
 	}
 
+	function deleteConnectionsByNodeId(
+		targetNodeId: string,
+		{ trackHistory = false, trackBulk = true } = {},
+	) {
+		const targetNode = workflowsStore.getNodeById(targetNodeId);
+		if (!targetNode) {
+			return;
+		}
+
+		if (trackHistory && trackBulk) {
+			historyStore.startRecordingUndo();
+		}
+
+		const connections = workflowsStore.workflow.connections;
+		for (const nodeName of Object.keys(connections)) {
+			const node = workflowsStore.getNodeByName(nodeName);
+			if (!node) {
+				continue;
+			}
+
+			for (const type of Object.keys(connections[nodeName])) {
+				for (const index of Object.keys(connections[nodeName][type])) {
+					for (const connectionIndex of Object.keys(
+						connections[nodeName][type][parseInt(index, 10)],
+					)) {
+						const connectionData =
+							connections[nodeName][type][parseInt(index, 10)][parseInt(connectionIndex, 10)];
+						if (!connectionData) {
+							continue;
+						}
+
+						const connectionDataNode = workflowsStore.getNodeByName(connectionData.node);
+						if (
+							connectionDataNode &&
+							(connectionDataNode.id === targetNode.id || node.name === targetNode.name)
+						) {
+							deleteConnection(
+								{
+									source: node.id,
+									sourceHandle: createCanvasConnectionHandleString({
+										mode: CanvasConnectionMode.Output,
+										type: type as NodeConnectionType,
+										index: parseInt(index, 10),
+									}),
+									target: connectionDataNode.id,
+									targetHandle: createCanvasConnectionHandleString({
+										mode: CanvasConnectionMode.Input,
+										type: connectionData.type as NodeConnectionType,
+										index: connectionData.index,
+									}),
+								},
+								{ trackHistory, trackBulk: false },
+							);
+						}
+					}
+				}
+			}
+		}
+
+		delete workflowsStore.workflow.connections[targetNode.name];
+
+		if (trackHistory && trackBulk) {
+			historyStore.stopRecordingUndo();
+		}
+	}
+
 	function deleteConnection(
 		connection: Connection,
 		{ trackHistory = false, trackBulk = true } = {},
@@ -1777,6 +1865,7 @@ export function useCanvasOperations({ router }: { router: ReturnType<typeof useR
 		revertCreateConnection,
 		deleteConnection,
 		revertDeleteConnection,
+		deleteConnectionsByNodeId,
 		isConnectionAllowed,
 		importWorkflowData,
 		fetchWorkflowDataFromUrl,
diff --git a/packages/editor-ui/src/stores/workflows.store.ts b/packages/editor-ui/src/stores/workflows.store.ts
index 02faa7ec4..c8464a5d4 100644
--- a/packages/editor-ui/src/stores/workflows.store.ts
+++ b/packages/editor-ui/src/stores/workflows.store.ts
@@ -1525,8 +1525,6 @@ export const useWorkflowsStore = defineStore(STORES.WORKFLOWS, () => {
 		}
 
 		removeNode(node);
-
-		// @TODO When removing node connected between two nodes, create a connection between them
 	}
 
 	function removeNodeConnectionsById(nodeId: string): void {