From e467c40c415e141da8a74c7b50de13768af05179 Mon Sep 17 00:00:00 2001
From: filtered <176114999+webfiltered@users.noreply.github.com>
Date: Sun, 11 May 2025 07:23:21 +1000
Subject: [PATCH] Fix DOM widgets after rebase

---
 src/scripts/app.ts       | 30 +++++++++++++++++++++++++++---
 src/scripts/domWidget.ts |  2 +-
 2 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/src/scripts/app.ts b/src/scripts/app.ts
index a5e0b9134..f4c8cefc6 100644
--- a/src/scripts/app.ts
+++ b/src/scripts/app.ts
@@ -5,7 +5,7 @@ import {
   LGraphNode,
   LiteGraph
 } from '@comfyorg/litegraph'
-import type { Vector2 } from '@comfyorg/litegraph'
+import type { Subgraph, Vector2 } from '@comfyorg/litegraph'
 import type { IBaseWidget } from '@comfyorg/litegraph/dist/types/widgets'
 import _ from 'lodash'
 import type { ToastMessageOptions } from 'primevue/toast'
@@ -43,6 +43,7 @@ import { useSubgraphService } from '@/services/subgraphService'
 import { useWorkflowService } from '@/services/workflowService'
 import { useApiKeyAuthStore } from '@/stores/apiKeyAuthStore'
 import { useCommandStore } from '@/stores/commandStore'
+import { useDomWidgetStore } from '@/stores/domWidgetStore'
 import { useExecutionStore } from '@/stores/executionStore'
 import { useExtensionStore } from '@/stores/extensionStore'
 import { useFirebaseAuthStore } from '@/stores/firebaseAuthStore'
@@ -72,6 +73,7 @@ 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,
@@ -715,7 +717,7 @@ export class ComfyApp {
   }
 
   #addAfterConfigureHandler() {
-    const { graph } = this
+    const { canvas, graph } = this
     const { onConfigure } = graph
     graph.onConfigure = function (...args) {
       fixLinkInputSlots(this)
@@ -732,7 +734,8 @@ export class ComfyApp {
         node.onAfterGraphConfigured?.()
       }
 
-      pruneWidgets(this.nodes)
+      const nodes = canvas.subgraph?.nodes ?? graph.nodes
+      pruneWidgets(nodes)
 
       return r
     }
@@ -791,6 +794,27 @@ export class ComfyApp {
     LiteGraph.alt_drag_do_clone_nodes = true
     LiteGraph.macGesturesRequireMac = false
 
+    this.canvas.canvas.addEventListener('litegraph:set-graph', (e) => {
+      type SetGraphType = CustomEvent<{
+        newGraph: LGraph | Subgraph
+        oldGraph: LGraph | Subgraph
+      }>
+
+      // Assertion: Not yet defined in litegraph.
+      const { newGraph } = (e as SetGraphType).detail
+
+      const nodeSet = new Set(newGraph.nodes)
+      const widgetStore = useDomWidgetStore()
+      for (const { widget } of widgetStore.widgetStates.values()) {
+        // Assertions: UnwrapRef
+        if (nodeSet.has(widget.node as LGraphNode)) {
+          widgetStore.registerWidget(widget as BaseDOMWidget)
+        } else {
+          widgetStore.unregisterWidget(widget.id)
+        }
+      }
+    })
+
     this.graph.start()
 
     // Ensure the canvas fills the window
diff --git a/src/scripts/domWidget.ts b/src/scripts/domWidget.ts
index 6acfcd95f..70f3bc31c 100644
--- a/src/scripts/domWidget.ts
+++ b/src/scripts/domWidget.ts
@@ -11,7 +11,7 @@ import type { InputSpec } from '@/schemas/nodeDef/nodeDefSchemaV2'
 import { useDomWidgetStore } from '@/stores/domWidgetStore'
 import { generateUUID } from '@/utils/formatUtil'
 
-export interface BaseDOMWidget<V extends object | string>
+export interface BaseDOMWidget<V extends object | string = object | string>
   extends IBaseWidget<V, string, DOMWidgetOptions<V>> {
   // ICustomWidget properties
   type: string