Clear DOM widgets instead of trying to manage refs

src/scripts/app.ts

@@ -5,7 +5,7 @@ import {
   LGraphNode,
   LiteGraph
 } from '@comfyorg/litegraph'
-import type { Subgraph, Vector2 } from '@comfyorg/litegraph'
+import type { Vector2 } from '@comfyorg/litegraph'
 import type { IBaseWidget } from '@comfyorg/litegraph/dist/types/widgets'
 import _ from 'lodash'
 import type { ToastMessageOptions } from 'primevue/toast'
@@ -74,8 +74,6 @@ import { deserialiseAndCreate } from '@/utils/vintageClipboard'
 
 import { type ComfyApi, PromptExecutionError, api } from './api'
 import { defaultGraph } from './defaultGraph'
-import type { BaseDOMWidget } from './domWidget'
-import { pruneWidgets } from './domWidget'
 import {
   getFlacMetadata,
   getLatentMetadata,
@@ -735,11 +733,6 @@ export class ComfyApp {
         node.onAfterGraphConfigured?.()
       }
 
-      graph.canvasAction((c) => {
-        const nodes = c.subgraph?.nodes ?? graph.nodes
-        pruneWidgets(nodes)
-      })
-
       return r
     }
   }
@@ -809,14 +802,14 @@ export class ComfyApp {
 
         // Assertions: UnwrapRef
         for (const { widget } of widgetStore.widgetStates.values()) {
-          if (!nodeSet.has(widget.node as LGraphNode)) {
-            widgetStore.unregisterWidget(widget.id)
+          if (!nodeSet.has(widget.node)) {
+            widgetStore.deactivateWidget(widget.id)
           }
         }
 
-        for (const { widget } of widgetStore.inactiveWidgetStates.values()) {
-          if (nodeSet.has(widget.node as LGraphNode)) {
-            widgetStore.registerWidget(widget as BaseDOMWidget)
+        for (const { widget } of widgetStore.inactiveWidgetStates) {
+          if (nodeSet.has(widget.node)) {
+            widgetStore.activateWidget(widget.id)
           }
         }
       }
@@ -1696,6 +1689,8 @@ export class ComfyApp {
     const executionStore = useExecutionStore()
     executionStore.lastNodeErrors = null
     executionStore.lastExecutionError = null
+
+    useDomWidgetStore().clear()
   }
 
   clientPosToCanvasPos(pos: Vector2): Vector2 {

src/services/workflowService.ts

@@ -7,6 +7,7 @@ import { ComfyWorkflowJSON } from '@/schemas/comfyWorkflowSchema'
 import { app } from '@/scripts/app'
 import { blankGraph, defaultGraph } from '@/scripts/defaultGraph'
 import { downloadBlob } from '@/scripts/utils'
+import { useDomWidgetStore } from '@/stores/domWidgetStore'
 import { useSettingStore } from '@/stores/settingStore'
 import { useToastStore } from '@/stores/toastStore'
 import { ComfyWorkflow, useWorkflowStore } from '@/stores/workflowStore'
@@ -20,6 +21,7 @@ export const useWorkflowService = () => {
   const workflowStore = useWorkflowStore()
   const toastStore = useToastStore()
   const dialogService = useDialogService()
+  const domWidgetStore = useDomWidgetStore()
 
   async function getFilename(defaultName: string): Promise<string | null> {
     if (settingStore.get('Comfy.PromptFilename')) {
@@ -285,11 +287,8 @@ export const useWorkflowService = () => {
    */
   const beforeLoadNewGraph = () => {
     // Use workspaceStore here as it is patched in unit tests.
-    const workflowStore = useWorkspaceStore().workflow
-    const activeWorkflow = workflowStore.activeWorkflow
-    if (activeWorkflow) {
-      activeWorkflow.changeTracker.store()
-    }
+    useWorkspaceStore().workflow.activeWorkflow?.changeTracker?.store()
+    domWidgetStore.clear()
   }
 
   /**

src/stores/domWidgetStore.ts

@@ -2,7 +2,7 @@
  * Stores all DOM widgets that are used in the canvas.
  */
 import { defineStore } from 'pinia'
-import { type Raw, markRaw, ref } from 'vue'
+import { type Raw, computed, markRaw, ref } from 'vue'
 
 import type { PositionConfig } from '@/composables/element/useAbsolutePosition'
 import type { BaseDOMWidget } from '@/scripts/domWidget'
@@ -13,11 +13,17 @@ export interface DomWidgetState extends PositionConfig {
   visible: boolean
   readonly: boolean
   zIndex: number
+  /** If the widget belongs to the current graph/subgraph. */
+  active: boolean
 }
 
 export const useDomWidgetStore = defineStore('domWidget', () => {
   const widgetStates = ref<Map<string, DomWidgetState>>(new Map())
 
+  const inactiveWidgetStates = computed(() =>
+    [...widgetStates.value.values()].filter((state) => !state.active)
+  )
+
   // Register a widget with the store
   const registerWidget = <V extends object | string>(
     widget: BaseDOMWidget<V>
@@ -28,7 +34,8 @@ export const useDomWidgetStore = defineStore('domWidget', () => {
       readonly: false,
       zIndex: 0,
       pos: [0, 0],
-      size: [0, 0]
+      size: [0, 0],
+      active: true
     })
   }
 
@@ -37,9 +44,27 @@ export const useDomWidgetStore = defineStore('domWidget', () => {
     widgetStates.value.delete(widgetId)
   }
 
+  const activateWidget = (widgetId: string) => {
+    const state = widgetStates.value.get(widgetId)
+    if (state) state.active = true
+  }
+
+  const deactivateWidget = (widgetId: string) => {
+    const state = widgetStates.value.get(widgetId)
+    if (state) state.active = false
+  }
+
+  const clear = () => {
+    widgetStates.value.clear()
+  }
+
   return {
     widgetStates,
+    inactiveWidgetStates,
     registerWidget,
-    unregisterWidget
+    unregisterWidget,
+    activateWidget,
+    deactivateWidget,
+    clear
   }
 })