[API] Improve widget typing (#1027)

2026-02-03 06:47:33 +00:00 · 2025-05-08 06:37:16 +10:00
parent 33a76714f7
commit 1b37502d70
12 changed files with 264 additions and 116 deletions
--- a/src/LGraphCanvas.ts
+++ b/src/LGraphCanvas.ts
@@ -34,7 +34,7 @@ import type {
  CanvasPointerExtensions,
 } from "./types/events"
 import type { ClipboardItems } from "./types/serialisation"
-import type { IWidget } from "./types/widgets"
+import type { IBaseWidget } from "./types/widgets"

 import { LinkConnector } from "@/canvas/LinkConnector"

@@ -466,7 +466,7 @@ export class LGraphCanvas {
  /** The current node being drawn by {@link drawNode}.  This should NOT be used to determine the currently selected node.  See {@link selectedItems} */
  current_node: LGraphNode | null
  /** used for widgets */
-  node_widget?: [LGraphNode, IWidget] | null
+  node_widget?: [LGraphNode, IBaseWidget] | null
  /** The link to draw a tooltip for. */
  over_link_center?: LinkSegment
  last_mouse_position: Point
@@ -1829,13 +1829,18 @@ export class LGraphCanvas {
  }

  /**
-   * Gets the widget at the current cursor position
+   * Gets the widget at the current cursor position.
   * @param node Optional node to check for widgets under cursor
-   * @returns The widget located at the current cursor position or null
+   * @returns The widget located at the current cursor position, if any is found.
+   * @deprecated Use {@link LGraphNode.getWidgetOnPos} instead.
+   * ```ts
+   * const [x, y] = canvas.graph_mouse
+   * const widget = canvas.node_over?.getWidgetOnPos(x, y, true)
+   * ```
   */
-  getWidgetAtCursor(node?: LGraphNode): IWidget | null {
+  getWidgetAtCursor(node?: LGraphNode): IBaseWidget | undefined {
    node ??= this.node_over
-    return node?.getWidgetOnPos(this.graph_mouse[0], this.graph_mouse[1], true) ?? null
+    return node?.getWidgetOnPos(this.graph_mouse[0], this.graph_mouse[1], true)
  }

  /**
@@ -1852,7 +1857,7 @@ export class LGraphCanvas {
    for (const otherNode of nodes) {
      if (otherNode.mouseOver && node != otherNode) {
        // mouse leave
-        otherNode.mouseOver = null
+        otherNode.mouseOver = undefined
        this._highlight_input = undefined
        this._highlight_pos = undefined
        this.linkConnector.overWidget = undefined
@@ -2381,7 +2386,7 @@ export class LGraphCanvas {
    this.dirty_canvas = true
  }

-  #processWidgetClick(e: CanvasPointerEvent, node: LGraphNode, widget: IWidget) {
+  #processWidgetClick(e: CanvasPointerEvent, node: LGraphNode, widget: IBaseWidget) {
    const { pointer } = this

    // Custom widget - CanvasPointer
@@ -2396,7 +2401,7 @@ export class LGraphCanvas {
    const x = pos[0] - node.pos[0]
    const y = pos[1] - node.pos[1]

-    const widgetInstance = toConcreteWidget(widget)
+    const widgetInstance = toConcreteWidget(widget, node, false)
    if (widgetInstance) {
      pointer.onClick = () => widgetInstance.onClick({
        e,
@@ -2629,15 +2634,11 @@ export class LGraphCanvas {
        const pos: Point = [0, 0]
        const inputId = isOverNodeInput(node, e.canvasX, e.canvasY, pos)
        const outputId = isOverNodeOutput(node, e.canvasX, e.canvasY, pos)
-        const overWidget = node.getWidgetOnPos(e.canvasX, e.canvasY, true)
+        const overWidget = node.getWidgetOnPos(e.canvasX, e.canvasY, true) ?? undefined

        if (!node.mouseOver) {
          // mouse enter
-          node.mouseOver = {
-            inputId: null,
-            outputId: null,
-            overWidget: null,
-          }
+          node.mouseOver = {}
          this.node_over = node
          this.dirty_canvas = true

@@ -2652,14 +2653,15 @@ export class LGraphCanvas {
        node.onMouseMove?.(e, [e.canvasX - node.pos[0], e.canvasY - node.pos[1]], this)

        // The input the mouse is over has changed
+        const { mouseOver } = node
        if (
-          node.mouseOver.inputId !== inputId ||
-          node.mouseOver.outputId !== outputId ||
-          node.mouseOver.overWidget !== overWidget
+          mouseOver.inputId !== inputId ||
+          mouseOver.outputId !== outputId ||
+          mouseOver.overWidget !== overWidget
        ) {
-          node.mouseOver.inputId = inputId
-          node.mouseOver.outputId = outputId
-          node.mouseOver.overWidget = overWidget
+          mouseOver.inputId = inputId
+          mouseOver.outputId = outputId
+          mouseOver.overWidget = overWidget

          // State reset
          linkConnector.overWidget = undefined
--- a/src/LGraphNode.ts
+++ b/src/LGraphNode.ts
@@ -24,7 +24,7 @@ import type { LGraph } from "./LGraph"
 import type { Reroute, RerouteId } from "./Reroute"
 import type { CanvasMouseEvent } from "./types/events"
 import type { ISerialisedNode } from "./types/serialisation"
-import type { IBaseWidget, IWidget, IWidgetOptions, TWidgetType, TWidgetValue } from "./types/widgets"
+import type { IBaseWidget, IWidgetOptions, TWidgetType, TWidgetValue } from "./types/widgets"

 import { getNodeInputOnPos, getNodeOutputOnPos } from "./canvas/measureSlots"
 import { NullGraphError } from "./infrastructure/NullGraphError"
@@ -46,7 +46,7 @@ import { findFreeSlotOfType } from "./utils/collections"
 import { distributeSpace } from "./utils/spaceDistribution"
 import { toClass } from "./utils/type"
 import { BaseWidget } from "./widgets/BaseWidget"
-import { WIDGET_TYPE_MAP } from "./widgets/widgetMap"
+import { toConcreteWidget, type WidgetTypeMap } from "./widgets/widgetMap"

 // #region Types

@@ -61,9 +61,9 @@ export interface INodePropertyInfo {
 }

 export interface IMouseOverData {
-  inputId: number | null
-  outputId: number | null
-  overWidget: IWidget | null
+  inputId?: number
+  outputId?: number
+  overWidget?: IBaseWidget
 }

 export interface ConnectByTypeOptions {
@@ -217,7 +217,7 @@ export class LGraphNode implements Positionable, IPinnable, IColorable {
  properties: Dictionary<NodeProperty | undefined> = {}
  properties_info: INodePropertyInfo[] = []
  flags: INodeFlags = {}
-  widgets?: IWidget[]
+  widgets?: IBaseWidget[]
  /**
   * The amount of space available for widgets to grow into.
   * @see {@link layoutWidgets}
@@ -324,7 +324,7 @@ export class LGraphNode implements Positionable, IPinnable, IColorable {
  console?: string[]
  _level?: number
  _shape?: RenderShape
-  mouseOver?: IMouseOverData | null
+  mouseOver?: IMouseOverData
  redraw_on_mouse?: boolean
  resizable?: boolean
  clonable?: boolean
@@ -510,7 +510,7 @@ export class LGraphNode implements Positionable, IPinnable, IColorable {
    name: string,
    value: unknown,
    old_value: unknown,
-    w: IWidget,
+    w: IBaseWidget,
  ): void
  onDeselected?(this: LGraphNode): void
  onKeyUp?(this: LGraphNode, e: KeyboardEvent): void
@@ -1675,13 +1675,13 @@ export class LGraphNode implements Positionable, IPinnable, IColorable {
   * @param options the object that contains special properties of this widget
   * @returns the created widget object
   */
-  addWidget(
-    type: TWidgetType,
+  addWidget<Type extends TWidgetType, TValue extends WidgetTypeMap[Type]["value"]>(
+    type: Type,
    name: string,
-    value: string | number | boolean | object,
-    callback: IWidget["callback"] | string | null,
+    value: TValue,
+    callback: IBaseWidget["callback"] | string | null,
    options?: IWidgetOptions | string,
-  ): IWidget {
+  ): WidgetTypeMap[Type] | IBaseWidget {
    this.widgets ||= []

    if (!options && callback && typeof callback === "object") {
@@ -1700,8 +1700,8 @@ export class LGraphNode implements Positionable, IPinnable, IColorable {
      callback = null
    }

-    const w: IWidget = {
-      // @ts-expect-error Type check or just assert?
+    const w: IBaseWidget & { type: Type } = {
+      // @ts-expect-error
      type: type.toLowerCase(),
      name: name,
      value: value,
@@ -1726,12 +1726,13 @@ export class LGraphNode implements Positionable, IPinnable, IColorable {
    return widget
  }

-  addCustomWidget<T extends IWidget>(custom_widget: T): T {
+  addCustomWidget<TPlainWidget extends IBaseWidget>(
+    custom_widget: TPlainWidget,
+  ): TPlainWidget | WidgetTypeMap[TPlainWidget["type"]] {
    this.widgets ||= []
-    const WidgetClass = WIDGET_TYPE_MAP[custom_widget.type]
-    const widget = WidgetClass ? new WidgetClass(custom_widget) as IWidget : custom_widget
+    const widget = toConcreteWidget(custom_widget, this, false) ?? custom_widget
    this.widgets.push(widget)
-    return widget as T
+    return widget
  }

  move(deltaX: number, deltaY: number): void {
@@ -1910,9 +1911,9 @@ export class LGraphNode implements Positionable, IPinnable, IColorable {
    canvasX: number,
    canvasY: number,
    includeDisabled = false,
-  ): IWidget | null {
+  ): IBaseWidget | undefined {
    const { widgets, pos, size } = this
-    if (!widgets?.length) return null
+    if (!widgets?.length) return

    const x = canvasX - pos[0]
    const y = canvasY - pos[1]
@@ -1938,7 +1939,6 @@ export class LGraphNode implements Positionable, IPinnable, IColorable {
        return widget
      }
    }
-    return null
  }

  /**
@@ -3368,7 +3368,7 @@ export class LGraphNode implements Positionable, IPinnable, IColorable {
  /**
   * Returns `true` if the widget is visible, otherwise `false`.
   */
-  isWidgetVisible(widget: IWidget): boolean {
+  isWidgetVisible(widget: IBaseWidget): boolean {
    const isHidden = (
      this.collapsed ||
      widget.hidden ||
@@ -3406,9 +3406,9 @@ export class LGraphNode implements Positionable, IPinnable, IColorable {
      if (widget.computedDisabled) ctx.globalAlpha *= 0.5
      const width = widget.width || nodeWidth

-      const WidgetClass: typeof WIDGET_TYPE_MAP[string] = WIDGET_TYPE_MAP[widget.type]
-      if (WidgetClass) {
-        toClass(WidgetClass, widget).drawWidget(ctx, { width, showText })
+      const widgetInstance = toConcreteWidget(widget, this, false)
+      if (widgetInstance) {
+        widgetInstance.drawWidget(ctx, { width, showText })
      } else {
        widget.draw?.(ctx, this, width, y, H, lowQuality)
      }
@@ -3482,7 +3482,7 @@ export class LGraphNode implements Positionable, IPinnable, IColorable {
    return this.#getMouseOverSlot(slot) === slot
  }

-  #isMouseOverWidget(widget: IWidget | undefined): boolean {
+  #isMouseOverWidget(widget: IBaseWidget | undefined): boolean {
    if (!widget) return false
    return this.mouseOver?.overWidget === widget
  }
@@ -3490,14 +3490,14 @@ export class LGraphNode implements Positionable, IPinnable, IColorable {
  /**
   * Returns the input slot that is associated with the given widget.
   */
-  getSlotFromWidget(widget: IWidget | undefined): INodeInputSlot | undefined {
+  getSlotFromWidget(widget: IBaseWidget | undefined): INodeInputSlot | undefined {
    if (widget) return this.inputs.find(slot => isWidgetInputSlot(slot) && slot.widget.name === widget.name)
  }

  /**
   * Returns the widget that is associated with the given input slot.
   */
-  getWidgetFromSlot(slot: INodeInputSlot): IWidget | undefined {
+  getWidgetFromSlot(slot: INodeInputSlot): IBaseWidget | undefined {
    if (!isWidgetInputSlot(slot)) return
    return this.widgets?.find(w => w.name === slot.widget.name)
  }
--- a/src/canvas/LinkConnector.ts
+++ b/src/canvas/LinkConnector.ts
@@ -4,7 +4,7 @@ import type { INodeInputSlot, INodeOutputSlot } from "@/interfaces"
 import type { LGraphNode } from "@/LGraphNode"
 import type { Reroute } from "@/Reroute"
 import type { CanvasPointerEvent } from "@/types/events"
-import type { IWidget } from "@/types/widgets"
+import type { IBaseWidget } from "@/types/widgets"

 import { CustomEventTarget } from "@/infrastructure/CustomEventTarget"
 import { LinkConnectorEventMap } from "@/infrastructure/LinkConnectorEventMap"
@@ -83,7 +83,7 @@ export class LinkConnector {
  readonly hiddenReroutes: Set<Reroute> = new Set()

  /** The widget beneath the pointer, if it is a valid connection target. */
-  overWidget?: IWidget
+  overWidget?: IBaseWidget
  /** The type (returned by downstream callback) for {@link overWidget} */
  overWidgetType?: string

--- a/src/litegraph.ts
+++ b/src/litegraph.ts
@@ -162,6 +162,8 @@ export { BooleanWidget } from "./widgets/BooleanWidget"
 export { ButtonWidget } from "./widgets/ButtonWidget"
 export { ComboWidget } from "./widgets/ComboWidget"
 export { KnobWidget } from "./widgets/KnobWidget"
+export { LegacyWidget } from "./widgets/LegacyWidget"
 export { NumberWidget } from "./widgets/NumberWidget"
 export { SliderWidget } from "./widgets/SliderWidget"
 export { TextWidget } from "./widgets/TextWidget"
+export { isComboWidget } from "./widgets/widgetMap"
--- a/src/types/widgets.ts
+++ b/src/types/widgets.ts
@@ -31,7 +31,7 @@ export interface IWidgetOptions<TValues = unknown[]> {
  callback?: IWidget["callback"]
 }

-export interface IWidgetSliderOptions extends IWidgetOptions<number> {
+export interface IWidgetSliderOptions extends IWidgetOptions<number[]> {
  min: number
  max: number
  step2: number
@@ -39,7 +39,7 @@ export interface IWidgetSliderOptions extends IWidgetOptions<number> {
  marker_color?: CanvasColour
 }

-export interface IWidgetKnobOptions extends IWidgetOptions<number> {
+export interface IWidgetKnobOptions extends IWidgetOptions<number[]> {
  min: number
  max: number
  step2: number
@@ -61,18 +61,19 @@ export type IWidget =
  | INumericWidget
  | IStringWidget
  | IComboWidget
+  | IStringComboWidget
  | ICustomWidget
  | ISliderWidget
  | IButtonWidget
  | IKnobWidget

-export interface IBooleanWidget extends IBaseWidget<boolean, "toggle", IWidgetOptions<boolean>> {
+export interface IBooleanWidget extends IBaseWidget<boolean, "toggle"> {
  type: "toggle"
  value: boolean
 }

 /** Any widget that uses a numeric backing */
-export interface INumericWidget extends IBaseWidget<number, "number", IWidgetOptions<number>> {
+export interface INumericWidget extends IBaseWidget<number, "number"> {
  type: "number"
  value: number
 }
@@ -89,6 +90,12 @@ export interface IKnobWidget extends IBaseWidget<number, "knob", IWidgetKnobOpti
  options: IWidgetKnobOptions
 }

+/** Avoids the type issues with the legacy IComboWidget type */
+export interface IStringComboWidget extends IBaseWidget<string, "combo", RequiredProps<IWidgetOptions<string[]>, "values">> {
+  type: "combo"
+  value: string
+}
+
 type ComboWidgetValues = string[] | Record<string, string> | ((widget?: IComboWidget, node?: LGraphNode) => string[])

 /** A combo-box widget (dropdown, select, etc) */
@@ -102,19 +109,19 @@ export interface IComboWidget extends IBaseWidget<
 }

 /** A widget with a string value */
-export interface IStringWidget extends IBaseWidget<string, "string" | "text", IWidgetOptions<string>> {
+export interface IStringWidget extends IBaseWidget<string, "string" | "text", IWidgetOptions<string[]>> {
  type: "string" | "text"
  value: string
 }

-export interface IButtonWidget extends IBaseWidget<undefined, "button", IWidgetOptions<undefined>> {
+export interface IButtonWidget extends IBaseWidget<string | undefined, "button"> {
  type: "button"
-  value: undefined
+  value: string | undefined
  clicked: boolean
 }

 /** A custom widget - accepts any value and has no built-in special handling */
-export interface ICustomWidget extends IBaseWidget<string | object, "custom", IWidgetOptions<string | object>> {
+export interface ICustomWidget extends IBaseWidget<string | object, "custom"> {
  type: "custom"
  value: string | object
 }
@@ -129,10 +136,17 @@ export type TWidgetValue = IWidget["value"]

 /**
 * The base type for all widgets.  Should not be implemented directly.
+ * @template TValue The type of value this widget holds.
+ * @template TType A string designating the type of widget, e.g. "toggle" or "string".
+ * @template TOptions The options for this widget.
 * @see IWidget
 */
-export interface IBaseWidget<TValue = unknown, TType = string, TOptions = unknown> {
-  linkedWidgets?: IWidget[]
+export interface IBaseWidget<
+  TValue = boolean | number | string | object | undefined,
+  TType extends string = string,
+  TOptions extends IWidgetOptions<unknown> = IWidgetOptions<unknown>,
+> {
+  linkedWidgets?: IBaseWidget[]

  name: string
  options: TOptions
--- a/src/widgets/BaseSteppedWidget.ts
+++ b/src/widgets/BaseSteppedWidget.ts
@@ -1,11 +1,11 @@
-import type { IWidget } from "@/types/widgets"
+import type { IBaseWidget } from "@/types/widgets"

 import { BaseWidget, type DrawWidgetOptions, type WidgetEventOptions } from "./BaseWidget"

 /**
 * Base class for widgets that have increment and decrement buttons.
 */
-export abstract class BaseSteppedWidget<TWidget extends IWidget> extends BaseWidget<TWidget> {
+export abstract class BaseSteppedWidget<TWidget extends IBaseWidget = IBaseWidget> extends BaseWidget<TWidget> {
  /**
   * Whether the widget can increment its value
   * @returns `true` if the widget can increment its value, otherwise `false`
--- a/src/widgets/BaseWidget.ts
+++ b/src/widgets/BaseWidget.ts
@@ -1,6 +1,6 @@
 import type { CanvasPointer, LGraphCanvas, LGraphNode, Size } from "@/litegraph"
 import type { CanvasMouseEvent, CanvasPointerEvent } from "@/types/events"
-import type { IBaseWidget, IWidget } from "@/types/widgets"
+import type { IBaseWidget } from "@/types/widgets"

 import { drawTextInArea } from "@/draw"
 import { Rectangle } from "@/infrastructure/Rectangle"
@@ -29,7 +29,7 @@ export interface WidgetEventOptions {
  canvas: LGraphCanvas
 }

-export abstract class BaseWidget<TWidget extends IWidget = IWidget> implements IBaseWidget {
+export abstract class BaseWidget<TWidget extends IBaseWidget = IBaseWidget> implements IBaseWidget {
  /** From node edge to widget edge */
  static margin = 15
  /** From widget edge to tip of arrow button */
@@ -41,12 +41,17 @@ export abstract class BaseWidget<TWidget extends IWidget = IWidget> implements I
  /** Minimum gap between label and value */
  static labelValueGap = 5

-  linkedWidgets?: IWidget[]
+  #node: LGraphNode
+  /** The node that this widget belongs to. */
+  get node() {
+    return this.#node
+  }
+
+  linkedWidgets?: IBaseWidget[]
  name: string
  options: TWidget["options"]
  label?: string
  type: TWidget["type"]
-  value: TWidget["value"]
  y: number = 0
  last_y?: number
  width?: number
@@ -64,27 +69,37 @@ export abstract class BaseWidget<TWidget extends IWidget = IWidget> implements I
    e?: CanvasMouseEvent,
  ): void
  mouse?(event: CanvasPointerEvent, pointerOffset: Point, node: LGraphNode): boolean
-  draw?(
-    ctx: CanvasRenderingContext2D,
-    node: LGraphNode,
-    widget_width: number,
-    y: number,
-    H: number,
-  ): void
  computeSize?(width?: number): Size
  onPointerDown?(pointer: CanvasPointer, node: LGraphNode, canvas: LGraphCanvas): boolean

-  constructor(widget: TWidget) {
+  #value: TWidget["value"]
+  get value(): TWidget["value"] {
+    return this.#value
+  }
+
+  set value(value: TWidget["value"]) {
+    this.#value = value
+  }
+
+  constructor(widget: TWidget & { node: LGraphNode })
+  constructor(widget: TWidget, node: LGraphNode)
+  constructor(widget: TWidget & { node: LGraphNode }, node?: LGraphNode) {
+    // Private fields
+    this.#node = node ?? widget.node
+    this.#value = widget.value
+
+    // `node` has no setter - Object.assign will throw.
    // TODO: Resolve this workaround. Ref: https://github.com/Comfy-Org/litegraph.js/issues/1022
    // @ts-expect-error Prevent naming conflicts with custom nodes.
    // eslint-disable-next-line unused-imports/no-unused-vars
-    const { outline_color, background_color, height, text_color, secondary_text_color, disabledTextColor, displayName, displayValue, labelBaseline, ...safeValues } = widget
+    const { node: _, outline_color, background_color, height, text_color, secondary_text_color, disabledTextColor, displayName, displayValue, labelBaseline, ...safeValues } = widget

    Object.assign(this, safeValues)
+
+    // Re-assign to fix TS errors.
    this.name = widget.name
    this.options = widget.options
    this.type = widget.type
-    this.value = widget.value
  }

  get outline_color() {
@@ -254,7 +269,7 @@ export abstract class BaseWidget<TWidget extends IWidget = IWidget> implements I
    const pos = canvas.graph_mouse
    this.callback?.(this.value, canvas, node, pos, e)

-    node.onWidgetChanged?.(this.name ?? "", v, oldValue, this as IWidget)
+    node.onWidgetChanged?.(this.name ?? "", v, oldValue, this)
    if (node.graph) node.graph._version++
  }
 }
--- a/src/widgets/ButtonWidget.ts
+++ b/src/widgets/ButtonWidget.ts
@@ -1,3 +1,4 @@
+import type { LGraphNode } from "@/LGraphNode"
 import type { IButtonWidget } from "@/types/widgets"

 import { BaseWidget, type DrawWidgetOptions, type WidgetEventOptions } from "./BaseWidget"
@@ -6,8 +7,8 @@ export class ButtonWidget extends BaseWidget<IButtonWidget> implements IButtonWi
  override type = "button" as const
  clicked: boolean

-  constructor(widget: IButtonWidget) {
-    super(widget)
+  constructor(widget: IButtonWidget, node: LGraphNode) {
+    super(widget, node)
    this.clicked ??= false
  }

--- a/src/widgets/ComboWidget.ts
+++ b/src/widgets/ComboWidget.ts
@@ -1,6 +1,6 @@
 import type { WidgetEventOptions } from "./BaseWidget"
 import type { LGraphNode } from "@/LGraphNode"
-import type { IComboWidget } from "@/types/widgets"
+import type { IComboWidget, IStringComboWidget } from "@/types/widgets"

 import { clamp, LiteGraph } from "@/litegraph"
 import { warnDeprecated } from "@/utils/feedback"
@@ -19,7 +19,7 @@ function toArray(values: Values): string[] {
  return Array.isArray(values) ? values : Object.keys(values)
 }

-export class ComboWidget extends BaseSteppedWidget<IComboWidget> implements IComboWidget {
+export class ComboWidget extends BaseSteppedWidget<IStringComboWidget | IComboWidget> implements IComboWidget {
  override type = "combo" as const

  override get _displayValue() {
--- a/src/widgets/LegacyWidget.ts
+++ b/src/widgets/LegacyWidget.ts
@@ -0,0 +1,33 @@
+import type { LGraphNode } from "@/LGraphNode"
+import type { IBaseWidget } from "@/types/widgets"
+
+import { LiteGraph } from "@/litegraph"
+
+import { BaseWidget, type DrawWidgetOptions } from "./BaseWidget"
+
+/**
+ * Wraps a legacy POJO custom widget, so that all widgets may be called via the same internal interface.
+ *
+ * Support will eventually be removed.
+ * @remarks Expect this class to undergo breaking changes without warning.
+ */
+export class LegacyWidget<TWidget extends IBaseWidget = IBaseWidget> extends BaseWidget<TWidget> implements IBaseWidget {
+  draw?(
+    ctx: CanvasRenderingContext2D,
+    node: LGraphNode,
+    widget_width: number,
+    y: number,
+    H: number,
+    lowQuality?: boolean,
+  ): void
+
+  override drawWidget(ctx: CanvasRenderingContext2D, options: DrawWidgetOptions) {
+    const H = LiteGraph.NODE_WIDGET_HEIGHT
+    const thisAsICustomWidget = this
+    thisAsICustomWidget.draw?.(ctx, this.node, options.width, this.y, H, !!options.showText)
+  }
+
+  override onClick() {
+    console.warn("Custom widget wrapper onClick was just called. Handling for third party widgets is done via LGraphCanvas - the mouse callback.")
+  }
+}
--- a/src/widgets/TextWidget.ts
+++ b/src/widgets/TextWidget.ts
@@ -1,10 +1,11 @@
+import type { LGraphNode } from "@/LGraphNode"
 import type { IStringWidget } from "@/types/widgets"

 import { BaseWidget, type DrawWidgetOptions, type WidgetEventOptions } from "./BaseWidget"

 export class TextWidget extends BaseWidget<IStringWidget> implements IStringWidget {
-  constructor(widget: IStringWidget) {
-    super(widget)
+  constructor(widget: IStringWidget, node: LGraphNode) {
+    super(widget, node)
    this.type ??= "string"
    this.value = widget.value?.toString() ?? ""
  }
--- a/src/widgets/widgetMap.ts
+++ b/src/widgets/widgetMap.ts
@@ -1,48 +1,128 @@
-import type { IBaseWidget, IWidget } from "@/types/widgets"
+import type { LGraphNode } from "@/LGraphNode"
+import type {
+  IBaseWidget,
+  IBooleanWidget,
+  IButtonWidget,
+  IComboWidget,
+  ICustomWidget,
+  IKnobWidget,
+  INumericWidget,
+  ISliderWidget,
+  IStringWidget,
+  IWidget,
+  TWidgetType,
+} from "@/types/widgets"
+
+import { toClass } from "@/utils/type"

 import { BaseWidget } from "./BaseWidget"
 import { BooleanWidget } from "./BooleanWidget"
 import { ButtonWidget } from "./ButtonWidget"
 import { ComboWidget } from "./ComboWidget"
 import { KnobWidget } from "./KnobWidget"
+import { LegacyWidget } from "./LegacyWidget"
 import { NumberWidget } from "./NumberWidget"
 import { SliderWidget } from "./SliderWidget"
 import { TextWidget } from "./TextWidget"

-export function toConcreteWidget(widget: IWidget): BaseWidget | undefined {
-  if (widget instanceof BaseWidget) return widget
+export type WidgetTypeMap = {
+  button: ButtonWidget
+  toggle: BooleanWidget
+  slider: SliderWidget
+  knob: KnobWidget
+  combo: ComboWidget
+  number: NumberWidget
+  string: TextWidget
+  text: TextWidget
+  custom: LegacyWidget
+  [key: string]: BaseWidget
+}

-  switch (widget.type) {
-  case "button": return new ButtonWidget(widget)
-  case "toggle": return new BooleanWidget(widget)
-  case "slider": return new SliderWidget(widget)
-  case "knob": return new KnobWidget(widget)
-  case "combo": return new ComboWidget(widget)
-  case "number": return new NumberWidget(widget)
-  case "string": return new TextWidget(widget)
-  case "text": return new TextWidget(widget)
+/**
+ * Convert a widget POJO to a proper widget instance.
+ * @param widget The POJO to convert.
+ * @param node The node the widget belongs to.
+ * @param wrapLegacyWidgets Whether to wrap legacy widgets in a `LegacyWidget` instance.
+ * @returns A concrete widget instance.
+ */
+export function toConcreteWidget<TWidget extends IWidget | IBaseWidget>(
+  widget: TWidget,
+  node: LGraphNode,
+  wrapLegacyWidgets?: true,
+): WidgetTypeMap[TWidget["type"]]
+export function toConcreteWidget<TWidget extends IWidget | IBaseWidget>(
+  widget: TWidget,
+  node: LGraphNode,
+  wrapLegacyWidgets: false): WidgetTypeMap[TWidget["type"]] | undefined
+export function toConcreteWidget<TWidget extends IWidget | IBaseWidget>(
+  widget: TWidget,
+  node: LGraphNode,
+  wrapLegacyWidgets = true,
+): WidgetTypeMap[TWidget["type"]] | undefined {
+  // Assertion: TypeScript has no concept of "all strings except X"
+  type RemoveBaseWidgetType<T> = T extends { type: TWidgetType } ? T : never
+  const narrowedWidget = widget as RemoveBaseWidgetType<TWidget>
+
+  switch (narrowedWidget.type) {
+  case "button": return toClass(ButtonWidget, narrowedWidget, node)
+  case "toggle": return toClass(BooleanWidget, narrowedWidget, node)
+  case "slider": return toClass(SliderWidget, narrowedWidget, node)
+  case "knob": return toClass(KnobWidget, narrowedWidget, node)
+  case "combo": return toClass(ComboWidget, narrowedWidget, node)
+  case "number": return toClass(NumberWidget, narrowedWidget, node)
+  case "string": return toClass(TextWidget, narrowedWidget, node)
+  case "text": return toClass(TextWidget, narrowedWidget, node)
+  default: {
+    if (wrapLegacyWidgets) return toClass(LegacyWidget, widget, node)
+  }
  }
 }

-type WidgetConstructor = {
-  new (plain: IBaseWidget): BaseWidget
+// #region Type Guards
+
+/** Type guard: Narrow **from {@link IBaseWidget}** to {@link IButtonWidget}. */
+export function isButtonWidget(widget: IBaseWidget): widget is IButtonWidget {
+  return widget.type === "button"
 }

-export const WIDGET_TYPE_MAP: Record<string, WidgetConstructor> = {
-  // @ts-expect-error https://github.com/Comfy-Org/litegraph.js/issues/616
-  button: ButtonWidget,
-  // @ts-expect-error #616
-  toggle: BooleanWidget,
-  // @ts-expect-error #616
-  slider: SliderWidget,
-  // @ts-expect-error #616
-  knob: KnobWidget,
-  // @ts-expect-error #616
-  combo: ComboWidget,
-  // @ts-expect-error #616
-  number: NumberWidget,
-  // @ts-expect-error #616
-  string: TextWidget,
-  // @ts-expect-error #616
-  text: TextWidget,
+/** Type guard: Narrow **from {@link IBaseWidget}** to {@link IBooleanWidget}. */
+export function isBooleanWidget(widget: IBaseWidget): widget is IBooleanWidget {
+  return widget.type === "toggle"
 }
+
+/** Type guard: Narrow **from {@link IBaseWidget}** to {@link ISliderWidget}. */
+export function isSliderWidget(widget: IBaseWidget): widget is ISliderWidget {
+  return widget.type === "slider"
+}
+
+/** Type guard: Narrow **from {@link IBaseWidget}** to {@link IKnobWidget}. */
+export function isKnobWidget(widget: IBaseWidget): widget is IKnobWidget {
+  return widget.type === "knob"
+}
+
+/** Type guard: Narrow **from {@link IBaseWidget}** to {@link IComboWidget}. */
+export function isComboWidget(widget: IBaseWidget): widget is IComboWidget {
+  return widget.type === "combo"
+}
+
+/** Type guard: Narrow **from {@link IBaseWidget}** to {@link INumericWidget}. */
+export function isNumberWidget(widget: IBaseWidget): widget is INumericWidget {
+  return widget.type === "number"
+}
+
+/** Type guard: Narrow **from {@link IBaseWidget}** to {@link IStringWidget}. */
+export function isStringWidget(widget: IBaseWidget): widget is IStringWidget {
+  return widget.type === "string"
+}
+
+/** Type guard: Narrow **from {@link IBaseWidget}** to {@link ITextWidget}. */
+export function isTextWidget(widget: IBaseWidget): widget is IStringWidget {
+  return widget.type === "text"
+}
+
+/** Type guard: Narrow **from {@link IBaseWidget}** to {@link ICustomWidget}. */
+export function isCustomWidget(widget: IBaseWidget): widget is ICustomWidget {
+  return widget.type === "custom"
+}
+
+// #endregion Type Guards