Retain workflow versions when serializing node_versions (#1563)

* retain initial node_versions on serialize

* give precedence to workflow version

* set version info on node

* move patch to setup hook

* switch to nodeCreated

js/workflow-metadata.js

@@ -3,12 +3,21 @@
  * - custom node pack version to all custom nodes used in the workflow
  *
  * Example metadata:
- "extra": {
+ * "nodes": {
-    "node_versions": {
+ *   "1": {
-      "comfy-core": "v0.3.8-4-g0b2eb7f",
+ *     type: "CheckpointLoaderSimple",
-      "comfyui-easy-use": "1.2.5"
+ *     ...
-    }
+ *     properties: {
-  },
+ *       cnr_id: "comfy-core",
+ *       version: "0.3.8",
+ *     },
+ *   },
+ * }
+ *
+ * @typedef {Object} NodeInfo
+ * @property {string} ver - Version (git hash or semantic version)
+ * @property {string} cnr_id - ComfyRegistry node ID
+ * @property {boolean} enabled - Whether the node is enabled
  */
 
 import { app } from "../../scripts/app.js";
@@ -23,7 +32,7 @@ class WorkflowMetadataExtension {
 
   /**
    * Get the installed nodes info
-   * @returns {Promise<Record<string, {ver: string, cnr_id: string, enabled: boolean}>>} The mapping from node name to its info.
+   * @returns {Promise<Record<string, NodeInfo>>} The mapping from node name to its info.
    * ver can either be a git commit hash or a semantic version such as "1.0.0"
    * cnr_id is the id of the node in the ComfyRegistry
    * enabled is true if the node is enabled, false if it is disabled
@@ -33,61 +42,41 @@ class WorkflowMetadataExtension {
     return await res.json();
   }
 
-  /**
-   * Get the node versions for the given graph
-   * @param {LGraph} graph The graph to get the node versions for
-   * @returns {Promise<Record<string, string>>} The mapping from node name to version
-   */
-  getGraphNodeVersions(graph) {
-    const nodeVersions = {};
-    for (const node of graph.nodes) {
-      const nodeData = node.constructor.nodeData;
-      // Frontend only nodes don't have nodeData
-      if (!nodeData) {
-        continue;
-      }
-      const modules = nodeData.python_module.split(".");
-
-      if (modules[0] === "custom_nodes") {
-        const nodePackageName = modules[1];
-        const nodeInfo =
-          this.installedNodes[nodePackageName] ??
-          this.installedNodes[nodePackageName.toLowerCase()];
-        if (nodeInfo) {
-          nodeVersions[nodePackageName] = nodeInfo.ver;
-        }
-      } else if (["nodes", "comfy_extras"].includes(modules[0])) {
-        nodeVersions["comfy-core"] = this.comfyCoreVersion;
-      } else {
-        console.warn(`Unknown node source: ${nodeData.python_module}`);
-      }
-    }
-    return nodeVersions;
-  }
-
   async init() {
-    const extension = this;
     this.installedNodes = await this.getInstalledNodes();
     this.comfyCoreVersion = (await api.getSystemStats()).system.comfyui_version;
+  }
 
-    // Attach metadata when app.graphToPrompt is called.
+  /**
-    const originalSerialize = LGraph.prototype.serialize;
+   * Called when any node is created
-    LGraph.prototype.serialize = function () {
+   * @param {LGraphNode} node The newly created node
-      const workflow = originalSerialize.apply(this, arguments);
+   */
+  nodeCreated(node) {
+    try {
+      // nodeData doesn't exist if node is missing or node is frontend only node
+      if (!node?.constructor?.nodeData?.python_module) return;
 
-      // Add metadata to the workflow
+      const nodeProperties = (node.properties ??= {});
-      if (!workflow.extra) {
+      const modules = node.constructor.nodeData.python_module.split(".");
-        workflow.extra = {};
+      const moduleType = modules[0];
+
+      if (moduleType === "custom_nodes") {
+        const nodePackageName = modules[1];
+        const { cnr_id, ver } =
+          this.installedNodes[nodePackageName] ??
+          this.installedNodes[nodePackageName.toLowerCase()] ??
+          {};
+
+        if (cnr_id === "comfy-core") return; // don't allow hijacking comfy-core name
+        if (cnr_id) nodeProperties.cnr_id = cnr_id;
+        if (ver) nodeProperties.ver = ver;
+      } else if (["nodes", "comfy_extras"].includes(moduleType)) {
+        nodeProperties.cnr_id = "comfy-core";
+        nodeProperties.ver = this.comfyCoreVersion;
       }
-      const graph = this;
+    } catch (e) {
-      try {
+      console.error(e);
-        workflow.extra["node_versions"] = extension.getGraphNodeVersions(graph);
+    }
-      } catch (e) {
-        console.error(e);
-      }
-
-      return workflow;
-    };
   }
 }