set version info on node

js/workflow-metadata.js

@@ -3,12 +3,21 @@
  * - custom node pack version to all custom nodes used in the workflow
  *
  * Example metadata:
- "extra": {
-    "node_versions": {
-      "comfy-core": "v0.3.8-4-g0b2eb7f",
-      "comfyui-easy-use": "1.2.5"
-    }
-  },
+ * "nodes": {
+ *   "1": {
+ *     type: "CheckpointLoaderSimple",
+ *     ...
+ *     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
@@ -34,44 +43,80 @@ class WorkflowMetadataExtension {
   }
 
   /**
-   * 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
+   * Set node versions for the given graph
+   * @param {LGraph} graph The graph to process
+   * @param {Object} workflow The serialized workflow object
    */
-  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(".");
+  setGraphNodeVersions(graph, workflow) {
+    if (!graph?.nodes || !workflow?.nodes) return;
 
-      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}`);
+    // Create a map of workflow nodes by ID
+    const workflowNodesById = {};
+    for (const node of workflow.nodes) {
+      if (node.id != null) {
+        workflowNodesById[node.id] = node;
       }
     }
-    return nodeVersions;
+
+    // Process each graph node and find its corresponding workflow node by ID
+    for (const graphNode of graph.nodes) {
+      const workflowNode = workflowNodesById[graphNode.id];
+      if (!workflowNode) continue;
+      this.setNodeVersion(graphNode, workflowNode);
+    }
   }
 
   /**
-   * The node versions embedded in the workflow json's initial state.
-   * @returns {Record<string, string> | undefined} The mapping from node name to version
+   * Set version information for a single node
+   * @param {Object} graphNode The graph node
+   * @param {Object} workflowNode The workflow node
    */
-  get workflowNodeVersions() {
-    return app.extensionManager?.workflow?.activeWorkflow?.initialState?.extra
-      ?.node_versions;
+  setNodeVersion(graphNode, workflowNode) {
+    const nodeProperties = (workflowNode.properties ??= {});
+    const nodeData = graphNode.constructor?.nodeData;
+
+    // The node is missing or is a frontend only node
+    // (node was not constructed in registerNodes closure where nodeData is set)
+    if (!nodeData) {
+      return;
+    }
+
+    const modules = nodeData.python_module.split(".");
+    const moduleType = modules[0];
+
+    if (moduleType === "custom_nodes") {
+      this.setCustomNodeVersion(modules[1], nodeProperties);
+    } else if (["nodes", "comfy_extras"].includes(moduleType)) {
+      this.setCoreNodeVersion(nodeProperties);
+    } else {
+      console.warn(`Unknown node source: ${nodeData.python_module}`);
+    }
+  }
+
+  /**
+   * Set version for custom nodes
+   * @private
+   */
+  setCustomNodeVersion(nodePackageName, nodeProperties) {
+    const nodeInfo =
+      this.installedNodes[nodePackageName] ??
+      this.installedNodes[nodePackageName.toLowerCase()];
+
+    if (nodeInfo) {
+      if (nodeInfo.cnr_id === "comfy-core") return; // reserved package name
+      // Preserve workflow cnr_id and version if they exist
+      nodeProperties.cnr_id ??= nodeInfo.cnr_id;
+      nodeProperties.version ??= nodeInfo.ver;
+    }
+  }
+
+  /**
+   * Set version for core nodes
+   * @private
+   */
+  setCoreNodeVersion(nodeProperties) {
+    nodeProperties.cnr_id ??= "comfy-core";
+    nodeProperties.version ??= this.comfyCoreVersion;
   }
 
   async init() {
@@ -84,20 +129,13 @@ class WorkflowMetadataExtension {
     LGraph.prototype.serialize = function () {
       const workflow = originalSerialize.apply(this, arguments);
 
-      // Add metadata to the workflow
-      if (!workflow.extra) {
-        workflow.extra = {};
-      }
+      // Add metadata to the nodes
       const graph = this;
       try {
-        workflow.extra["node_versions"] = {
-          ...extension.getGraphNodeVersions(graph),
-          ...extension.workflowNodeVersions, // give precedence to the workflow node versions
-        };
+        extension.setGraphNodeVersions(graph, workflow);
       } catch (e) {
         console.error(e);
       }
-
       return workflow;
     };
   }