[feat] Update node output system to use NodeLocatorIds (#4500)

Author: christian-byrne

src/extensions/core/uploadAudio.ts

@@ -1,6 +1,7 @@
 import type { LGraphNode } from '@comfyorg/litegraph'
 import type { IStringWidget } from '@comfyorg/litegraph/dist/types/widgets'
 
+import { useChainCallback } from '@/composables/functional/useChainCallback'
 import { useNodeDragAndDrop } from '@/composables/node/useNodeDragAndDrop'
 import { useNodeFileInput } from '@/composables/node/useNodeFileInput'
 import { useNodePaste } from '@/composables/node/useNodePaste'
@@ -9,6 +10,8 @@ import type { ResultItemType } from '@/schemas/apiSchema'
 import type { ComfyNodeDef } from '@/schemas/nodeDefSchema'
 import type { DOMWidget } from '@/scripts/domWidget'
 import { useToastStore } from '@/stores/toastStore'
+import { NodeLocatorId } from '@/types'
+import { getNodeByLocatorId } from '@/utils/graphTraversalUtil'
 
 import { api } from '../../scripts/api'
 import { app } from '../../scripts/app'
@@ -137,14 +140,27 @@ app.registerExtension({
             audioUIWidget.element.classList.remove('empty-audio-widget')
           }
         }
+
+        audioUIWidget.onRemove = useChainCallback(
+          audioUIWidget.onRemove,
+          () => {
+            if (!audioUIWidget.element) return
+            audioUIWidget.element.pause()
+            audioUIWidget.element.src = ''
+            audioUIWidget.element.remove()
+          }
+        )
+
         return { widget: audioUIWidget }
       }
     }
   },
-  onNodeOutputsUpdated(nodeOutputs: Record<number, any>) {
-    for (const [nodeId, output] of Object.entries(nodeOutputs)) {
-      const node = app.graph.getNodeById(nodeId)
+  onNodeOutputsUpdated(nodeOutputs: Record<NodeLocatorId, any>) {
+    for (const [nodeLocatorId, output] of Object.entries(nodeOutputs)) {
       if ('audio' in output) {
+        const node = getNodeByLocatorId(app.graph, nodeLocatorId)
+        if (!node) continue
+
         // @ts-expect-error fixme ts strict error
         const audioUIWidget = node.widgets.find(
           (w) => w.name === 'audioUI'

src/scripts/app.ts

@@ -47,6 +47,7 @@ import { useDomWidgetStore } from '@/stores/domWidgetStore'
 import { useExecutionStore } from '@/stores/executionStore'
 import { useExtensionStore } from '@/stores/extensionStore'
 import { useFirebaseAuthStore } from '@/stores/firebaseAuthStore'
+import { useNodeOutputStore } from '@/stores/imagePreviewStore'
 import { KeyComboImpl, useKeybindingStore } from '@/stores/keybindingStore'
 import { useModelStore } from '@/stores/modelStore'
 import { SYSTEM_NODE_DEFS, useNodeDefStore } from '@/stores/nodeDefStore'
@@ -60,6 +61,10 @@ import type { ComfyExtension, MissingNodeType } from '@/types/comfy'
 import { ExtensionManager } from '@/types/extensionTypes'
 import { ColorAdjustOptions, adjustColor } from '@/utils/colorUtil'
 import { graphToPrompt } from '@/utils/executionUtil'
+import {
+  getNodeByExecutionId,
+  triggerCallbackOnAllNodes
+} from '@/utils/graphTraversalUtil'
 import {
   executeWidgetsCallback,
   fixLinkInputSlots,
@@ -640,29 +645,21 @@ export class ComfyApp {
     })
 
     api.addEventListener('executed', ({ detail }) => {
-      const output = this.nodeOutputs[detail.display_node || detail.node]
-      if (detail.merge && output) {
-        for (const k in detail.output ?? {}) {
-          const v = output[k]
-          if (v instanceof Array) {
-            output[k] = v.concat(detail.output[k])
-          } else {
-            output[k] = detail.output[k]
-          }
-        }
-      } else {
-        this.nodeOutputs[detail.display_node || detail.node] = detail.output
-      }
-      const node = this.graph.getNodeById(detail.display_node || detail.node)
-      if (node) {
-        if (node.onExecuted) node.onExecuted(detail.output)
+      const nodeOutputStore = useNodeOutputStore()
+      const executionId = String(detail.display_node || detail.node)
+
+      nodeOutputStore.setNodeOutputsByExecutionId(executionId, detail.output, {
+        merge: detail.merge
+      })
+
+      const node = getNodeByExecutionId(this.graph, executionId)
+      if (node && node.onExecuted) {
+        node.onExecuted(detail.output)
       }
     })
 
     api.addEventListener('execution_start', () => {
-      this.graph.nodes.forEach((node) => {
-        if (node.onExecutionStart) node.onExecutionStart()
-      })
+      triggerCallbackOnAllNodes(this.graph, 'onExecutionStart')
     })
 
     api.addEventListener('execution_error', ({ detail }) => {
@@ -690,11 +687,13 @@ export class ComfyApp {
     api.addEventListener('b_preview_with_metadata', ({ detail }) => {
       // Enhanced preview with explicit node context
       const { blob, displayNodeId } = detail
+      const { setNodePreviewsByExecutionId, revokePreviewsByExecutionId } =
+        useNodeOutputStore()
       // Ensure clean up if `executing` event is missed.
-      this.revokePreviews(displayNodeId)
+      revokePreviewsByExecutionId(displayNodeId)
       const blobUrl = URL.createObjectURL(blob)
-      // Preview cleanup is now handled in progress_state event to support multiple concurrent previews
-      this.nodePreviewImages[displayNodeId] = [blobUrl]
+      // Preview cleanup is handled in progress_state event to support multiple concurrent previews
+      setNodePreviewsByExecutionId(displayNodeId, [blobUrl])
     })
 
     api.init()
@@ -1673,25 +1672,13 @@ export class ComfyApp {
     }
   }
 
-  /**
-   * Frees memory allocated to image preview blobs for a specific node, by revoking the URLs associated with them.
-   * @param nodeId ID of the node to revoke all preview images of
-   */
-  revokePreviews(nodeId: NodeId) {
-    if (!this.nodePreviewImages[nodeId]?.[Symbol.iterator]) return
-    for (const url of this.nodePreviewImages[nodeId]) {
-      URL.revokeObjectURL(url)
-    }
-  }
   /**
    * Clean current state
    */
   clean() {
     this.nodeOutputs = {}
-    for (const id of Object.keys(this.nodePreviewImages)) {
-      this.revokePreviews(id)
-    }
-    this.nodePreviewImages = {}
+    const { revokeAllPreviews } = useNodeOutputStore()
+    revokeAllPreviews()
     const executionStore = useExecutionStore()
     executionStore.lastNodeErrors = null
     executionStore.lastExecutionError = null

src/stores/executionStore.ts

@@ -24,6 +24,7 @@ import type {
 } from '@/schemas/comfyWorkflowSchema'
 import { api } from '@/scripts/api'
 import { app } from '@/scripts/app'
+import { useNodeOutputStore } from '@/stores/imagePreviewStore'
 import type { NodeLocatorId } from '@/types/nodeIdentification'
 import { createNodeLocatorId } from '@/types/nodeIdentification'
 
@@ -229,9 +230,9 @@ export const useExecutionStore = defineStore('execution', () => {
     api.addEventListener('progress_state', handleProgressState)
     api.addEventListener('status', handleStatus)
     api.addEventListener('execution_error', handleExecutionError)
+    api.addEventListener('progress_text', handleProgressText)
+    api.addEventListener('display_component', handleDisplayComponent)
   }
-  api.addEventListener('progress_text', handleProgressText)
-  api.addEventListener('display_component', handleDisplayComponent)
 
   function unbindExecutionEvents() {
     api.removeEventListener('execution_start', handleExecutionStart)
@@ -244,6 +245,7 @@ export const useExecutionStore = defineStore('execution', () => {
     api.removeEventListener('status', handleStatus)
     api.removeEventListener('execution_error', handleExecutionError)
     api.removeEventListener('progress_text', handleProgressText)
+    api.removeEventListener('display_component', handleDisplayComponent)
   }
 
   function handleExecutionStart(e: CustomEvent<ExecutionStartWsMessage>) {
@@ -294,8 +296,8 @@ export const useExecutionStore = defineStore('execution', () => {
         // Note that we're doing the *actual* node id instead of the display node id
         // here intentionally. That way, we don't clear the preview every time a new node
         // within an expanded graph starts executing.
-        app.revokePreviews(nodeId)
-        delete app.nodePreviewImages[nodeId]
+        const { revokePreviewsByExecutionId } = useNodeOutputStore()
+        revokePreviewsByExecutionId(nodeId)
       }
     }