Fix change tracker count desync on error (#1555)

* Add TS types

* Ensure changeTracker works after exceptions

Wraps all code between before/after change calls in try/finally blocks

src/extensions/core/groupNode.ts

@@ -13,6 +13,7 @@ import {
   deserialiseAndCreate,
   serialise
 } from '@/extensions/core/vintageClipboard'
+import type { ComfyNodeDef } from '@/types/apiTypes'
 
 type GroupNodeWorkflowData = {
   external: ComfyLink[]
@@ -56,7 +57,7 @@ const Workflow = {
 
 class GroupNodeBuilder {
   nodes: LGraphNode[]
-  nodeData: any
+  nodeData: GroupNodeWorkflowData
 
   constructor(nodes: LGraphNode[]) {
     this.nodes = nodes
@@ -175,7 +176,7 @@ export class GroupNodeConfig {
   primitiveToWidget: {}
   nodeInputs: {}
   outputVisibility: any[]
-  nodeDef: any
+  nodeDef: ComfyNodeDef
   inputs: any[]
   linksFrom: {}
   linksTo: {}
@@ -204,6 +205,7 @@ export class GroupNodeConfig {
       output: [],
       output_name: [],
       output_is_list: [],
+      // @ts-expect-error Unused, doesn't exist
       output_is_hidden: [],
       name: source + SEPARATOR + this.name,
       display_name: this.name,
@@ -695,11 +697,11 @@ export class GroupNodeConfig {
 }
 
 export class GroupNodeHandler {
-  node
+  node: LGraphNode
   groupData
   innerNodes: any
 
-  constructor(node) {
+  constructor(node: LGraphNode) {
     this.node = node
     this.groupData = node.constructor?.nodeData?.[GROUP]
 
@@ -774,6 +776,7 @@ export class GroupNodeHandler {
 
     this.node.updateLink = (link) => {
       // Replace the group node reference with the internal node
+      // @ts-expect-error Can this be removed?  Or replaced with: LLink.create(link.asSerialisable())
       link = { ...link }
       const output = this.groupData.newToOldOutputMap[link.origin_slot]
       let innerNode = this.innerNodes[output.node.index]
@@ -965,17 +968,20 @@ export class GroupNodeHandler {
 
       app.canvas.emitBeforeChange()
 
-      const { newNodes, selectedIds } = addInnerNodes()
-      reconnectInputs(selectedIds)
-      reconnectOutputs(selectedIds)
-      app.graph.remove(this.node)
+      try {
+        const { newNodes, selectedIds } = addInnerNodes()
+        reconnectInputs(selectedIds)
+        reconnectOutputs(selectedIds)
+        app.graph.remove(this.node)
 
-      app.canvas.emitAfterChange()
-
-      return newNodes
+        return newNodes
+      } finally {
+        app.canvas.emitAfterChange()
+      }
     }
 
     const getExtraMenuOptions = this.node.getExtraMenuOptions
+    // @ts-expect-error Should pass patched return value getExtraMenuOptions
     this.node.getExtraMenuOptions = function (_, options) {
       getExtraMenuOptions?.apply(this, arguments)
 
@@ -988,6 +994,7 @@ export class GroupNodeHandler {
         null,
         {
           content: 'Convert to nodes',
+          // @ts-expect-error
           callback: () => {
             return this.convertToNodes()
           }
@@ -1148,6 +1155,7 @@ export class GroupNodeHandler {
 
           if (
             old.inputName !== 'image' &&
+            // @ts-expect-error Widget values
             !widget.options.values.includes(widget.value)
           ) {
             widget.value = widget.options.values[0]
@@ -1354,6 +1362,7 @@ export class GroupNodeHandler {
       if (!originNode) continue // this node is in the group
       originNode.connect(
         originSlot,
+        // @ts-expect-error Valid - uses deprecated interface.  Required check: if (graph.getNodeById(this.node.id) !== this.node) report()
         this.node.id,
         this.groupData.oldToNewInputMap[targetId][targetSlot]
       )
@@ -1475,7 +1484,7 @@ function ungroupSelectedGroupNodes() {
   const nodes = Object.values(app.canvas.selected_nodes ?? {})
   for (const node of nodes) {
     if (GroupNodeHandler.isGroupNode(node)) {
-      node['convertToNodes']?.()
+      node.convertToNodes?.()
     }
   }
 }

src/extensions/core/vintageClipboard.ts

@@ -69,51 +69,54 @@ export function deserialiseAndCreate(data: string, canvas: LGraphCanvas): void {
 
   const { graph, graph_mouse } = canvas
   canvas.emitBeforeChange()
-  graph.beforeChange()
+  try {
+    graph.beforeChange()
 
-  const deserialised = JSON.parse(data)
+    const deserialised = JSON.parse(data)
 
-  // Find the top left point of the boundary of all pasted nodes
-  const topLeft = [Infinity, Infinity]
-  for (const { pos } of deserialised.nodes) {
-    if (topLeft[0] > pos[0]) topLeft[0] = pos[0]
-    if (topLeft[1] > pos[1]) topLeft[1] = pos[1]
+    // Find the top left point of the boundary of all pasted nodes
+    const topLeft = [Infinity, Infinity]
+    for (const { pos } of deserialised.nodes) {
+      if (topLeft[0] > pos[0]) topLeft[0] = pos[0]
+      if (topLeft[1] > pos[1]) topLeft[1] = pos[1]
+    }
+
+    // Silent default instead of throw
+    if (!Number.isFinite(topLeft[0]) || !Number.isFinite(topLeft[1])) {
+      topLeft[0] = graph_mouse[0]
+      topLeft[1] = graph_mouse[1]
+    }
+
+    // Create nodes
+    const nodes: LGraphNode[] = []
+    for (const info of deserialised.nodes) {
+      const node = LiteGraph.createNode(info.type)
+      if (!node) continue
+
+      node.configure(info)
+
+      // Paste to the bottom right of pointer
+      node.pos[0] += graph_mouse[0] - topLeft[0]
+      node.pos[1] += graph_mouse[1] - topLeft[1]
+
+      graph.add(node, true)
+      nodes.push(node)
+    }
+
+    // Create links
+    for (const info of deserialised.links) {
+      const relativeId = info[0]
+      const outNode = relativeId != null ? nodes[relativeId] : undefined
+
+      const inNode = nodes[info[2]]
+      if (outNode && inNode) outNode.connect(info[1], inNode, info[3])
+      else console.warn('Warning, nodes missing on pasting')
+    }
+
+    canvas.selectNodes(nodes)
+
+    graph.afterChange()
+  } finally {
+    canvas.emitAfterChange()
   }
-
-  // Silent default instead of throw
-  if (!Number.isFinite(topLeft[0]) || !Number.isFinite(topLeft[1])) {
-    topLeft[0] = graph_mouse[0]
-    topLeft[1] = graph_mouse[1]
-  }
-
-  // Create nodes
-  const nodes: LGraphNode[] = []
-  for (const info of deserialised.nodes) {
-    const node = LiteGraph.createNode(info.type)
-    if (!node) continue
-
-    node.configure(info)
-
-    // Paste to the bottom right of pointer
-    node.pos[0] += graph_mouse[0] - topLeft[0]
-    node.pos[1] += graph_mouse[1] - topLeft[1]
-
-    graph.add(node, true)
-    nodes.push(node)
-  }
-
-  // Create links
-  for (const info of deserialised.links) {
-    const relativeId = info[0]
-    const outNode = relativeId != null ? nodes[relativeId] : undefined
-
-    const inNode = nodes[info[2]]
-    if (outNode && inNode) outNode.connect(info[1], inNode, info[3])
-    else console.warn('Warning, nodes missing on pasting')
-  }
-
-  canvas.selectNodes(nodes)
-
-  graph.afterChange()
-  canvas.emitAfterChange()
 }

src/scripts/app.ts

@@ -1573,10 +1573,7 @@ export class ComfyApp {
     api.addEventListener('execution_start', ({ detail }) => {
       this.lastExecutionError = null
       this.graph.nodes.forEach((node) => {
-        // @ts-expect-error
-        if (node.onExecutionStart)
-          // @ts-expect-error
-          node.onExecutionStart()
+        if (node.onExecutionStart) node.onExecutionStart()
       })
     })
 
@@ -2408,8 +2405,8 @@ export class ComfyApp {
         }
       }
 
-      const innerNodes = outerNode['getInnerNodes']
-        ? outerNode['getInnerNodes']()
+      const innerNodes = outerNode.getInnerNodes
+        ? outerNode.getInnerNodes()
         : [outerNode]
       for (const node of innerNodes) {
         if (node.isVirtualNode) {
@@ -2427,8 +2424,8 @@ export class ComfyApp {
     for (const outerNode of graph.computeExecutionOrder(false)) {
       const skipNode = outerNode.mode === 2 || outerNode.mode === 4
       const innerNodes =
-        !skipNode && outerNode['getInnerNodes']
-          ? outerNode['getInnerNodes']()
+        !skipNode && outerNode.getInnerNodes
+          ? outerNode.getInnerNodes()
           : [outerNode]
       for (const node of innerNodes) {
         if (node.isVirtualNode) {
@@ -2894,7 +2891,6 @@ export class ComfyApp {
     for (let nodeNum in this.graph.nodes) {
       const node = this.graph.nodes[nodeNum]
       const def = defs[node.type]
-      // @ts-expect-error
       // Allow primitive nodes to handle refresh
       node.refreshComboInNode?.(defs)
 

src/types/litegraph-augmentation.d.ts

@@ -1,6 +1,7 @@
 import '@comfyorg/litegraph'
 import type { ComfyNodeDef } from '@/types/apiTypes'
 import type { LLink } from '@comfyorg/litegraph'
+import type { NodeId } from './comfyWorkflow'
 
 /**
  *  ComfyUI extensions of litegraph
@@ -26,8 +27,17 @@ declare module '@comfyorg/litegraph' {
     onExecuted?(output: any): void
     onNodeCreated?(this: LGraphNode): void
     setInnerNodes?(nodes: LGraphNode[]): void
+    // TODO: Requires several coercion changes to runtime code.
+    getInnerNodes?() // : LGraphNode[]
+    convertToNodes?(): LGraphNode[]
+    recreate?(): Promise<LGraphNode>
+    refreshComboInNode?(defs: Record<string, ComfyNodeDef>)
     applyToGraph?(extraLinks?: LLink[]): void
     updateLink?(link: LLink): LLink | null
+    onExecutionStart?(): unknown
+
+    index?: number
+    runningInternalNodeId?: NodeId
 
     comfyClass?: string