ComfyUI_frontend/src/extensions/core/proxyWidget.ts

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 { parseProxyWidgets } from '@/schemas/proxyWidget'
import { DOMWidgetImpl } from '@/scripts/domWidget'
import { useExtensionService } from '@/services/extensionService'
import { useDomWidgetStore } from '@/stores/domWidgetStore'
import { useCanvasStore } from '@/stores/graphStore'

useExtensionService().registerExtension({
  name: 'Comfy.SubgraphProxyWidgets',
  nodeCreated(node: LGraphNode) {
    if (node instanceof SubgraphNode) {
      setTimeout(() => injectProperty(node), 0)
    }
  }
})
function injectProperty(subgraphNode: SubgraphNode) {
  const canvasStore = useCanvasStore()
  subgraphNode.properties.proxyWidgets ??= []
  const proxyWidgets = subgraphNode.properties.proxyWidgets
  Object.defineProperty(subgraphNode.properties, 'proxyWidgets', {
    get: () => {
      const result = subgraphNode.widgets
        .filter((w) => isProxyWidget(w))
        .map((w) => [w._overlay.nodeId, w._overlay.widgetName])
      return JSON.stringify(result)
    },
    set: (property: string) => {
      const parsed = parseProxyWidgets(property)
      const { widgetStates } = useDomWidgetStore()
      for (const w of subgraphNode.widgets ?? []) {
        if (w instanceof DOMWidgetImpl && widgetStates.has(w.id)) {
          const widgetState = widgetStates.get(w.id)
          if (!widgetState) continue
          widgetState.active = false
        }
      }
      //NOTE: This does not apply to pushed entries, only initial load
      subgraphNode.widgets = subgraphNode.widgets.filter(
        (w) => !isProxyWidget(w)
      )
      for (const [nodeId, widgetName] of parsed) {
        const w = addProxyWidget(subgraphNode, `${nodeId}`, widgetName)
        if (w instanceof DOMWidgetImpl) {
          const widgetState = widgetStates.get(w.id)
          if (!widgetState) continue
          widgetState.active = true
          widgetState.widget = w
        }
      }
      canvasStore.canvas?.setDirty(true, true)
    }
  })
  subgraphNode.properties.proxyWidgets = proxyWidgets
}
type Overlay = Partial<IBaseWidget> & {
  graph: LGraph
  nodeId: string
  widgetName: string
  isProxyWidget: boolean
}
type ProxyWidget = IBaseWidget & { _overlay: Overlay }
function isProxyWidget(w: IBaseWidget): w is ProxyWidget {
  return (w as { _overlay?: Overlay })?._overlay?.isProxyWidget ?? false
}

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,
    lasy_y: 0,
    width: undefined,
    computedHeight: 0,
    afterQueued: undefined,
    onRemove: undefined,
    node: subgraphNode
  }
  return addProxyFromOverlay(subgraphNode, overlay)
}
function resolveLinkedWidget(
  overlay: Overlay
): [LGraphNode | undefined, IBaseWidget | undefined] {
  const { graph, nodeId, widgetName } = overlay
  let g: LGraph | undefined = graph
  let n: LGraphNode | SubgraphNode | undefined = undefined
  for (const id of nodeId.split(':')) {
    n = g?._nodes_by_id?.[id]
    g = n?.isSubgraphNode?.() ? n.subgraph : undefined
  }
  if (!n) return [undefined, undefined]
  return [n, n.widgets?.find((w: IBaseWidget) => w.name === widgetName)]
}
function addProxyFromOverlay(subgraphNode: SubgraphNode, overlay: Overlay) {
  //TODO: call toConcrete when resolved and hold reference?
  //NOTE: From testing, WeakRefs never dropped. May refactor later
  let [linkedNode, linkedWidget] = resolveLinkedWidget(overlay)
  const handler = Object.fromEntries(
    ['get', 'set', 'getPrototypeOf', 'ownKeys', 'has'].map((s) => {
      const func = function (t: object, p: string, ...rest: object[]) {
        if (s == 'get' && p == '_overlay') return overlay
        const bw = linkedWidget ?? disconnectedWidget
        if (s == 'set' && p == 'computedDisabled') {
          //ignore setting, calc actual
          bw.computedDisabled =
            bw.disabled || linkedNode!.getSlotFromWidget(bw)?.link != null
          return true
        }
        //NOTE: p may be undefined
        let r = rest.at(-1)
        if (overlay.hasOwnProperty(p)) r = t = overlay
        else {
          t = bw
          if (p == 'value') r = t
        }
        return (Reflect as any)[s](t, p, ...rest.slice(0, -1), r)
      }
      return [s, func]
    })
  )
  const w = new Proxy(overlay, handler) as IBaseWidget
  subgraphNode.widgets.push(w)
  return w
}