Subgraph widget promotion - Part 2 (#5617)

Implements proxyWidget support on subgraph nodes.  This registers a
special proxyWidgets property on subgraph nodes which is directly mapped
to the proxyWidgets displayed on the node. Each proxyWidget directly
maps to a real widget inside the subgraph.

┆Issue is synchronized with this [Notion
page](https://www.notion.so/PR-5617-Subgraph-widget-promotion-Part-2-2716d73d3650813d8621fefdce6ae518)
by [Unito](https://www.unito.io)

---------

Co-authored-by: GitHub Action <[email protected]>

Author: 2 people

knip.config.ts

@@ -25,9 +25,7 @@ const config: KnipConfig = {
     'src/workbench/extensions/manager/types/generatedManagerTypes.ts',
     'src/types/comfyRegistryTypes.ts',
     // Used by a custom node (that should move off of this)
-    'src/scripts/ui/components/splitButton.ts',
-    // Staged for for use with subgraph widget promotion
-    'src/lib/litegraph/src/widgets/DisconnectedWidget.ts'
+    'src/scripts/ui/components/splitButton.ts'
   ],
   compilers: {
     // https://github.com/webpro-nl/knip/issues/1008#issuecomment-3207756199

src/core/graph/subgraph/proxyWidget.ts

@@ -0,0 +1,185 @@
+import { useNodeImage } from '@/composables/node/useNodeImage'
+import { parseProxyWidgets } from '@/core/schemas/proxyWidget'
+import type { LGraph, LGraphNode } from '@/lib/litegraph/src/litegraph'
+import { SubgraphNode } from '@/lib/litegraph/src/subgraph/SubgraphNode'
+import type { IBaseWidget } from '@/lib/litegraph/src/types/widgets.ts'
+import { disconnectedWidget } from '@/lib/litegraph/src/widgets/DisconnectedWidget'
+import { useCanvasStore } from '@/renderer/core/canvas/canvasStore'
+import { DOMWidgetImpl } from '@/scripts/domWidget'
+import { useDomWidgetStore } from '@/stores/domWidgetStore'
+import { useNodeOutputStore } from '@/stores/imagePreviewStore'
+import { getNodeByExecutionId } from '@/utils/graphTraversalUtil'
+
+/**
+ * @typedef {object} Overlay - Each proxy Widget has an associated overlay object
+ * Accessing a property which exists in the overlay object will
+ * instead result in the action being performed on the overlay object
+ * 3 properties are added for locating the proxied widget
+ * @property {LGraph} graph - The graph the widget resides in. Used for widget lookup
+ * @property {string} nodeId - The NodeId the proxy Widget is located on
+ * @property {string} widgetName - The name of the linked widget
+ *
+ * @property {boolean} isProxyWidget - Always true, used as type guard
+ * @property {LGraphNode} node - not included on IBaseWidget, but required for overlay
+ */
+type Overlay = Partial<IBaseWidget> & {
+  graph: LGraph
+  nodeId: string
+  widgetName: string
+  isProxyWidget: boolean
+  node?: LGraphNode
+}
+// A ProxyWidget can be treated like a normal widget.
+// the _overlay property can be used to directly access the Overlay object
+/**
+ * @typedef {object} ProxyWidget - a reference to a widget that can
+ * be displayed and owned by a separate node
+ * @property {Overlay} _overlay - a special property to access the overlay of the widget
+ * Any property that exists in the overlay will be accessed instead of the property
+ * on the linked widget
+ */
+type ProxyWidget = IBaseWidget & { _overlay: Overlay }
+function isProxyWidget(w: IBaseWidget): w is ProxyWidget {
+  return (w as { _overlay?: Overlay })?._overlay?.isProxyWidget ?? false
+}
+
+const originalOnConfigure = SubgraphNode.prototype.onConfigure
+SubgraphNode.prototype.onConfigure = function (serialisedNode) {
+  if (!this.isSubgraphNode())
+    throw new Error("Can't add proxyWidgets to non-subgraphNode")
+
+  const canvasStore = useCanvasStore()
+  //Must give value to proxyWidgets prior to defining or it won't serialize
+  this.properties.proxyWidgets ??= '[]'
+  let proxyWidgets = this.properties.proxyWidgets
+
+  originalOnConfigure?.call(this, serialisedNode)
+
+  Object.defineProperty(this.properties, 'proxyWidgets', {
+    get: () => {
+      return proxyWidgets
+    },
+    set: (property: string) => {
+      const parsed = parseProxyWidgets(property)
+      const { deactivateWidget, setWidget } = useDomWidgetStore()
+      for (const w of this.widgets.filter((w) => isProxyWidget(w))) {
+        if (w instanceof DOMWidgetImpl) deactivateWidget(w.id)
+      }
+      this.widgets = this.widgets.filter((w) => !isProxyWidget(w))
+      for (const [nodeId, widgetName] of parsed) {
+        const w = addProxyWidget(this, `${nodeId}`, widgetName)
+        if (w instanceof DOMWidgetImpl) setWidget(w)
+      }
+      proxyWidgets = property
+      canvasStore.canvas?.setDirty(true, true)
+      this._setConcreteSlots()
+      this.arrange()
+    }
+  })
+  this.properties.proxyWidgets = proxyWidgets
+}
+
+function addProxyWidget(
+  subgraphNode: SubgraphNode,
+  nodeId: string,
+  widgetName: string
+) {
+  const name = `${nodeId}: ${widgetName}`
+  const overlay = {
+    nodeId,
+    widgetName,
+    graph: subgraphNode.subgraph,
+    name,
+    label: name,
+    isProxyWidget: true,
+    y: 0,
+    last_y: undefined,
+    width: undefined,
+    computedHeight: undefined,
+    afterQueued: undefined,
+    onRemove: undefined,
+    node: subgraphNode
+  }
+  return addProxyFromOverlay(subgraphNode, overlay)
+}
+function resolveLinkedWidget(
+  overlay: Overlay
+): [LGraphNode | undefined, IBaseWidget | undefined] {
+  const { graph, nodeId, widgetName } = overlay
+  const n = getNodeByExecutionId(graph, nodeId)
+  if (!n) return [undefined, undefined]
+  return [n, n.widgets?.find((w: IBaseWidget) => w.name === widgetName)]
+}
+function addProxyFromOverlay(subgraphNode: SubgraphNode, overlay: Overlay) {
+  let [linkedNode, linkedWidget] = resolveLinkedWidget(overlay)
+  let backingWidget = linkedWidget ?? disconnectedWidget
+  if (overlay.widgetName == '$$canvas-image-preview')
+    overlay.node = new Proxy(subgraphNode, {
+      get(_t, p) {
+        if (p !== 'imgs') return Reflect.get(subgraphNode, p)
+        if (!linkedNode) return []
+        const images =
+          useNodeOutputStore().getNodeOutputs(linkedNode)?.images ?? []
+        if (images !== linkedNode.images) {
+          linkedNode.images = images
+          useNodeImage(linkedNode).showPreview()
+        }
+        return linkedNode.imgs
+      }
+    })
+  /**
+   * A set of handlers which define widget interaction
+   * Many arguments are shared between function calls
+   * @param {IBaseWidget} _t - The "target" the call is originally made on.
+   *   This argument is never used, but must be defined for typechecking
+   * @param {string} property - The name of the accessed value.
+   *   Checked for conditional logic, but never changed
+   * @param {object} receiver - The object the result is set to
+   *   and the vlaue used as 'this' if property is a get/set method
+   * @param {unknown} value - only used on set calls. The thing being assigned
+   */
+  const handler = {
+    get(_t: IBaseWidget, property: string, receiver: object) {
+      let redirectedTarget: object = backingWidget
+      let redirectedReceiver = receiver
+      if (property == '_overlay') return overlay
+      else if (property == 'value') redirectedReceiver = backingWidget
+      if (Object.prototype.hasOwnProperty.call(overlay, property)) {
+        redirectedTarget = overlay
+        redirectedReceiver = overlay
+      }
+      return Reflect.get(redirectedTarget, property, redirectedReceiver)
+    },
+    set(_t: IBaseWidget, property: string, value: unknown, receiver: object) {
+      let redirectedTarget: object = backingWidget
+      let redirectedReceiver = receiver
+      if (property == 'value') redirectedReceiver = backingWidget
+      else if (property == 'computedHeight') {
+        //update linkage regularly, but no more than once per frame
+        ;[linkedNode, linkedWidget] = resolveLinkedWidget(overlay)
+        backingWidget = linkedWidget ?? disconnectedWidget
+      }
+      if (Object.prototype.hasOwnProperty.call(overlay, property)) {
+        redirectedTarget = overlay
+        redirectedReceiver = overlay
+      }
+      return Reflect.set(redirectedTarget, property, value, redirectedReceiver)
+    },
+    getPrototypeOf() {
+      return Reflect.getPrototypeOf(backingWidget)
+    },
+    ownKeys() {
+      return Reflect.ownKeys(backingWidget)
+    },
+    has(_t: IBaseWidget, property: string) {
+      let redirectedTarget: object = backingWidget
+      if (Object.prototype.hasOwnProperty.call(overlay, property)) {
+        redirectedTarget = overlay
+      }
+      return Reflect.has(redirectedTarget, property)
+    }
+  }
+  const w = new Proxy(disconnectedWidget, handler)
+  subgraphNode.widgets.push(w)
+  return w
+}

src/core/schemas/proxyWidget.ts

@@ -0,0 +1,23 @@
+import { z } from 'zod'
+import { fromZodError } from 'zod-validation-error'
+
+import type { NodeProperty } from '@/lib/litegraph/src/LGraphNode'
+
+const proxyWidgetsPropertySchema = z.array(z.tuple([z.string(), z.string()]))
+type ProxyWidgetsProperty = z.infer<typeof proxyWidgetsPropertySchema>
+
+export function parseProxyWidgets(
+  property: NodeProperty | undefined
+): ProxyWidgetsProperty {
+  if (typeof property !== 'string') {
+    throw new Error(
+      'Invalid assignment for properties.proxyWidgets:\nValue must be a string'
+    )
+  }
+  const parsed = JSON.parse(property)
+  const result = proxyWidgetsPropertySchema.safeParse(parsed)
+  if (result.success) return result.data
+
+  const error = fromZodError(result.error)
+  throw new Error(`Invalid assignment for properties.proxyWidgets:\n${error}`)
+}

src/scripts/widgets.ts

@@ -1,3 +1,4 @@
+import '@/core/graph/subgraph/proxyWidget'
 import { t } from '@/i18n'
 import { type LGraphNode, isComboWidget } from '@/lib/litegraph/src/litegraph'
 import type {

src/stores/domWidgetStore.ts

@@ -57,6 +57,13 @@ export const useDomWidgetStore = defineStore('domWidget', () => {
     if (state) state.active = false
   }
 
+  const setWidget = (widget: BaseDOMWidget) => {
+    const state = widgetStates.value.get(widget.id)
+    if (!state) return
+    state.active = true
+    state.widget = widget
+  }
+
   const clear = () => {
     widgetStates.value.clear()
   }
@@ -69,6 +76,7 @@ export const useDomWidgetStore = defineStore('domWidget', () => {
     unregisterWidget,
     activateWidget,
     deactivateWidget,
+    setWidget,
     clear
   }
 })

tests-ui/tests/widgets/proxyWidget.test.ts

@@ -0,0 +1,119 @@
+import { describe, expect, test, vi } from 'vitest'
+
+import '@/core/graph/subgraph/proxyWidget'
+//import { ComponentWidgetImpl, DOMWidgetImpl } from '@/scripts/domWidget'
+
+import { LGraphNode, type SubgraphNode } from '@/lib/litegraph/src/litegraph'
+
+import {
+  createTestSubgraph,
+  createTestSubgraphNode
+} from '../litegraph/subgraph/fixtures/subgraphHelpers'
+
+vi.mock('@/renderer/core/canvas/canvasStore', () => ({
+  useCanvasStore: () => ({})
+}))
+vi.mock('@/stores/domWidgetStore', () => ({
+  useDomWidgetStore: () => ({ widgetStates: new Map() })
+}))
+
+function setupSubgraph(
+  innerNodeCount: number = 0
+): [SubgraphNode, LGraphNode[]] {
+  const subgraph = createTestSubgraph()
+  const subgraphNode = createTestSubgraphNode(subgraph)
+  subgraphNode._internalConfigureAfterSlots()
+  const graph = subgraphNode.graph
+  graph.add(subgraphNode)
+  const innerNodes = []
+  for (let i = 0; i < innerNodeCount; i++) {
+    const innerNode = new LGraphNode(`InnerNode${i}`)
+    subgraph.add(innerNode)
+    innerNodes.push(innerNode)
+  }
+  return [subgraphNode, innerNodes]
+}
+
+describe('Subgraph proxyWidgets', () => {
+  test('Can add simple widget', () => {
+    const [subgraphNode, innerNodes] = setupSubgraph(1)
+    innerNodes[0].addWidget('text', 'stringWidget', 'value', () => {})
+    subgraphNode.properties.proxyWidgets = JSON.stringify([
+      ['1', 'stringWidget']
+    ])
+    expect(subgraphNode.widgets.length).toBe(1)
+    expect(subgraphNode.properties.proxyWidgets).toBe(
+      JSON.stringify([['1', 'stringWidget']])
+    )
+  })
+  test('Can add multiple widgets with same name', () => {
+    const [subgraphNode, innerNodes] = setupSubgraph(2)
+    for (const innerNode of innerNodes)
+      innerNode.addWidget('text', 'stringWidget', 'value', () => {})
+    subgraphNode.properties.proxyWidgets = JSON.stringify([
+      ['1', 'stringWidget'],
+      ['2', 'stringWidget']
+    ])
+    expect(subgraphNode.widgets.length).toBe(2)
+    expect(subgraphNode.widgets[0].name).not.toEqual(
+      subgraphNode.widgets[1].name
+    )
+  })
+  test('Will not modify existing widgets', () => {
+    const [subgraphNode, innerNodes] = setupSubgraph(1)
+    innerNodes[0].addWidget('text', 'stringWidget', 'value', () => {})
+    subgraphNode.addWidget('text', 'stringWidget', 'value', () => {})
+    subgraphNode.properties.proxyWidgets = JSON.stringify([
+      ['1', 'stringWidget']
+    ])
+    expect(subgraphNode.widgets.length).toBe(2)
+    subgraphNode.properties.proxyWidgets = JSON.stringify([])
+    expect(subgraphNode.widgets.length).toBe(1)
+  })
+  test('Will mirror changes to value', () => {
+    const [subgraphNode, innerNodes] = setupSubgraph(1)
+    innerNodes[0].addWidget('text', 'stringWidget', 'value', () => {})
+    subgraphNode.properties.proxyWidgets = JSON.stringify([
+      ['1', 'stringWidget']
+    ])
+    expect(subgraphNode.widgets.length).toBe(1)
+    expect(subgraphNode.widgets[0].value).toBe('value')
+    innerNodes[0].widgets![0].value = 'test'
+    expect(subgraphNode.widgets[0].value).toBe('test')
+    subgraphNode.widgets[0].value = 'test2'
+    expect(innerNodes[0].widgets![0].value).toBe('test2')
+  })
+  test('Will not modify position or sizing of existing widgets', () => {
+    const [subgraphNode, innerNodes] = setupSubgraph(1)
+    innerNodes[0].addWidget('text', 'stringWidget', 'value', () => {})
+    subgraphNode.properties.proxyWidgets = JSON.stringify([
+      ['1', 'stringWidget']
+    ])
+    if (!innerNodes[0].widgets) throw new Error('node has no widgets')
+    innerNodes[0].widgets[0].y = 10
+    innerNodes[0].widgets[0].last_y = 11
+    innerNodes[0].widgets[0].computedHeight = 12
+    subgraphNode.widgets[0].y = 20
+    subgraphNode.widgets[0].last_y = 21
+    subgraphNode.widgets[0].computedHeight = 22
+    expect(innerNodes[0].widgets[0].y).toBe(10)
+    expect(innerNodes[0].widgets[0].last_y).toBe(11)
+    expect(innerNodes[0].widgets[0].computedHeight).toBe(12)
+  })
+  test('Can detatch and re-attach widgets', () => {
+    const [subgraphNode, innerNodes] = setupSubgraph(1)
+    innerNodes[0].addWidget('text', 'stringWidget', 'value', () => {})
+    subgraphNode.properties.proxyWidgets = JSON.stringify([
+      ['1', 'stringWidget']
+    ])
+    if (!innerNodes[0].widgets) throw new Error('node has no widgets')
+    expect(subgraphNode.widgets[0].value).toBe('value')
+    const poppedWidget = innerNodes[0].widgets.pop()
+    //simulate new draw frame
+    subgraphNode.widgets[0].computedHeight = 10
+    expect(subgraphNode.widgets[0].value).toBe(undefined)
+    innerNodes[0].widgets.push(poppedWidget!)
+    subgraphNode.widgets[0].computedHeight = 10
+    expect(subgraphNode.widgets[0].value).toBe('value')
+  })
+})