From f9ff5d910508e669a59e00aa79025d4de8ab4464 Mon Sep 17 00:00:00 2001
From: pythongosssss <125205205+pythongosssss@users.noreply.github.com>
Date: Sat, 4 Nov 2023 11:44:01 +0000
Subject: [PATCH] Handle missing internal nodes

---
 tests-ui/tests/groupNode.test.js | 39 ++++++++++++++++++++++++++++++++
 web/extensions/core/groupNode.js | 25 +++++++++++++++-----
 web/scripts/app.js               |  2 +-
 3 files changed, 59 insertions(+), 7 deletions(-)

diff --git a/tests-ui/tests/groupNode.test.js b/tests-ui/tests/groupNode.test.js
index 76890a76d..afa4b92b9 100644
--- a/tests-ui/tests/groupNode.test.js
+++ b/tests-ui/tests/groupNode.test.js
@@ -541,4 +541,43 @@ describe("group node", () => {
 			})
 		);
 	});
+	test("shows missing node error on missing internal node when loading graph data", async () => {
+		const { graph } = await start();
+
+		const dialogShow = jest.spyOn(graph.app.ui.dialog, "show");
+		await graph.app.loadGraphData({
+			last_node_id: 3,
+			last_link_id: 1,
+			nodes: [
+				{
+					id: 3,
+					type: "workflow/testerror",
+				},
+			],
+			links: [],
+			groups: [],
+			config: {},
+			extra: {
+				groupNodes: {
+					testerror: {
+						nodes: [
+							{
+								type: "NotKSampler",
+							},
+							{
+								type: "NotVAEDecode",
+							},
+						],
+					},
+				},
+			},
+		});
+
+		expect(dialogShow).toBeCalledTimes(1);
+		const call = dialogShow.mock.calls[0][0];
+		expect(call).toContain("the following node types were not found");
+		expect(call).toContain("NotKSampler");
+		expect(call).toContain("NotVAEDecode");
+		expect(call).toContain("workflow/testerror");
+	});
 });
diff --git a/web/extensions/core/groupNode.js b/web/extensions/core/groupNode.js
index c54d75979..d2655580a 100644
--- a/web/extensions/core/groupNode.js
+++ b/web/extensions/core/groupNode.js
@@ -6,7 +6,7 @@ export const IS_GROUP_NODE = Symbol();
 export const GROUP_DATA = Symbol();
 const GROUP_SLOTS = Symbol();
 
-export async function registerGroupNodes(groupNodes, source, prefix) {
+export async function registerGroupNodes(groupNodes, source, prefix, missingNodeTypes) {
 	if (!groupNodes) return;
 
 	let extra = app.graph.extra;
@@ -15,7 +15,20 @@ export async function registerGroupNodes(groupNodes, source, prefix) {
 	if (!nodes) extra.groupNodes = nodes = {};
 
 	for (const g in groupNodes) {
-		const def = buildNodeDef(groupNodes[g], g, globalDefs, source);
+		const groupData = groupNodes[g];
+
+		let hasMissing = false;
+		for (const n of groupData.nodes) {
+			// Find missing node types
+			if (!(n.type in LiteGraph.registered_node_types)) {
+				missingNodeTypes.push(n.type);
+				hasMissing = true;
+			}
+		}
+
+		if (hasMissing) continue;
+
+		const def = buildNodeDef(groupData, g, globalDefs, source);
 		if (prefix) {
 			def.display_name = prefix + "/" + def.display_name;
 		}
@@ -319,8 +332,8 @@ class ConvertToGroupAction {
 				Object.values(app.canvas.selected_nodes).find((n) => n.constructor.nodeData?.[IS_GROUP_NODE]),
 			callback: async () => {
 				const name = this.getName();
-				if(!name) return;
-				
+				if (!name) return;
+
 				let extra = app.graph.extra;
 				if (!extra) app.graph.extra = extra = {};
 				let groupNodes = extra.groupNodes;
@@ -402,8 +415,8 @@ const ext = {
 			};
 		}
 	},
-	async beforeConfigureGraph(graphData) {
-		registerGroupNodes(graphData?.extra?.groupNodes, "workflow");
+	async beforeConfigureGraph(graphData, missingNodeTypes) {
+		registerGroupNodes(graphData?.extra?.groupNodes, "workflow", undefined, missingNodeTypes);
 	},
 	addCustomNodeDefs(defs) {
 		globalDefs = defs;
diff --git a/web/scripts/app.js b/web/scripts/app.js
index ebbb64be6..73aea3afe 100644
--- a/web/scripts/app.js
+++ b/web/scripts/app.js
@@ -1486,8 +1486,8 @@ export class ComfyApp {
 			reset_invalid_values = true;
 		}
 
-		await this.#invokeExtensionsAsync("beforeConfigureGraph", graphData);
 		const missingNodeTypes = [];
+		await this.#invokeExtensionsAsync("beforeConfigureGraph", graphData, missingNodeTypes);
 		for (let n of graphData.nodes) {
 			// Patch T2IAdapterLoader to ControlNetLoader since they are the same node now
 			if (n.type == "T2IAdapterLoader") n.type = "ControlNetLoader";