[Schema] ComfyNodeDefV2 schema (#2847)

src/components/graph/NodeTooltip.vue

@@ -76,7 +76,7 @@ const onIdle = () => {
     const inputName = node.inputs[inputSlot].name
     const translatedTooltip = st(
       `nodeDefs.${normalizeI18nKey(node.type)}.inputs.${normalizeI18nKey(inputName)}.tooltip`,
-      nodeDef.inputs.getInput(inputName)?.tooltip
+      nodeDef.inputs[inputName]?.tooltip
     )
     return showTooltip(translatedTooltip)
   }
@@ -90,7 +90,7 @@ const onIdle = () => {
   if (outputSlot !== -1) {
     const translatedTooltip = st(
       `nodeDefs.${normalizeI18nKey(node.type)}.outputs.${outputSlot}.tooltip`,
-      nodeDef.outputs.all?.[outputSlot]?.tooltip
+      nodeDef.outputs[outputSlot]?.tooltip
     )
     return showTooltip(translatedTooltip)
   }
@@ -100,7 +100,7 @@ const onIdle = () => {
   if (widget && !widget.element) {
     const translatedTooltip = st(
       `nodeDefs.${normalizeI18nKey(node.type)}.inputs.${normalizeI18nKey(widget.name)}.tooltip`,
-      nodeDef.inputs.getInput(widget.name)?.tooltip
+      nodeDef.inputs[widget.name]?.tooltip
     )
     // Widget tooltip can be set dynamically, current translation collection does not support this.
     return showTooltip(widget.tooltip ?? translatedTooltip)

src/components/node/NodePreview.vue

@@ -83,16 +83,13 @@ https://github.com/Nuked88/ComfyUI-N-Sidebar/blob/7ae7da4a9761009fb6629bc04c6830
 import _ from 'lodash'
 import { computed } from 'vue'
 
-import { ComfyNodeDefImpl } from '@/stores/nodeDefStore'
+import type { ComfyNodeDef as ComfyNodeDefV2 } from '@/schemas/nodeDef/nodeDefSchemaV2'
 import { useWidgetStore } from '@/stores/widgetStore'
 import { useColorPaletteStore } from '@/stores/workspace/colorPaletteStore'
 
-const props = defineProps({
+const props = defineProps<{
-  nodeDef: {
+  nodeDef: ComfyNodeDefV2
-    type: ComfyNodeDefImpl,
+}>()
-    required: true
-  }
-})
 
 const colorPaletteStore = useColorPaletteStore()
 const litegraphColors = computed(
@@ -102,8 +99,8 @@ const litegraphColors = computed(
 const widgetStore = useWidgetStore()
 
 const nodeDef = props.nodeDef
-const allInputDefs = nodeDef.inputs.all
+const allInputDefs = Object.values(nodeDef.inputs)
-const allOutputDefs = nodeDef.outputs.all
+const allOutputDefs = nodeDef.outputs
 const slotInputDefs = allInputDefs.filter(
   (input) => !widgetStore.inputIsWidget(input)
 )

src/extensions/core/widgetInputs.ts

@@ -624,7 +624,7 @@ app.registerExtension({
     app.canvas.getWidgetLinkType = function (widget, node) {
       const nodeDefStore = useNodeDefStore()
       const nodeDef = nodeDefStore.nodeDefsByName[node.type]
-      const input = nodeDef.inputs.getInput(widget.name)
+      const input = nodeDef.inputs[widget.name]
       return input?.type
     }
 

src/schemas/nodeDef/migration.ts

@@ -0,0 +1,121 @@
+import {
+  ComfyNodeDef as ComfyNodeDefV2,
+  InputSpec as InputSpecV2,
+  OutputSpec as OutputSpecV2
+} from '@/schemas/nodeDef/nodeDefSchemaV2'
+import {
+  ComfyNodeDef as ComfyNodeDefV1,
+  InputSpec as InputSpecV1,
+  getComboSpecComboOptions,
+  isComboInputSpec,
+  isComboInputSpecV1
+} from '@/schemas/nodeDefSchema'
+
+/**
+ * Transforms a V1 node definition to V2 format
+ * @param nodeDefV1 The V1 node definition to transform
+ * @returns The transformed V2 node definition
+ */
+export function transformNodeDefV1ToV2(
+  nodeDefV1: ComfyNodeDefV1
+): ComfyNodeDefV2 {
+  // Transform inputs
+  const inputs: Record<string, InputSpecV2> = {}
+
+  // Process required inputs
+  if (nodeDefV1.input?.required) {
+    Object.entries(nodeDefV1.input.required).forEach(([name, inputSpecV1]) => {
+      inputs[name] = transformInputSpec(inputSpecV1, name, false)
+    })
+  }
+
+  // Process optional inputs
+  if (nodeDefV1.input?.optional) {
+    Object.entries(nodeDefV1.input.optional).forEach(([name, inputSpecV1]) => {
+      inputs[name] = transformInputSpec(inputSpecV1, name, true)
+    })
+  }
+
+  // Transform outputs
+  const outputs: OutputSpecV2[] = []
+
+  if (nodeDefV1.output) {
+    nodeDefV1.output.forEach((outputType, index) => {
+      const outputSpec: OutputSpecV2 = {
+        index,
+        name: nodeDefV1.output_name?.[index] || `output_${index}`,
+        type: Array.isArray(outputType) ? 'COMBO' : outputType,
+        is_list: nodeDefV1.output_is_list?.[index] || false,
+        tooltip: nodeDefV1.output_tooltips?.[index]
+      }
+
+      // Add options for combo outputs
+      if (Array.isArray(outputType)) {
+        outputSpec.options = outputType
+      }
+
+      outputs.push(outputSpec)
+    })
+  }
+
+  // Create the V2 node definition
+  return {
+    inputs,
+    outputs,
+    hidden: nodeDefV1.input?.hidden,
+    name: nodeDefV1.name,
+    display_name: nodeDefV1.display_name,
+    description: nodeDefV1.description,
+    category: nodeDefV1.category,
+    output_node: nodeDefV1.output_node,
+    python_module: nodeDefV1.python_module,
+    deprecated: nodeDefV1.deprecated,
+    experimental: nodeDefV1.experimental
+  }
+}
+
+/**
+ * Transforms a V1 input specification to V2 format
+ * @param inputSpecV1 The V1 input specification to transform
+ * @param name The name of the input
+ * @param isOptional Whether the input is optional
+ * @returns The transformed V2 input specification
+ */
+function transformInputSpec(
+  inputSpecV1: InputSpecV1,
+  name: string,
+  isOptional: boolean
+): InputSpecV2 {
+  // Extract options from the input spec
+  const options = inputSpecV1[1] || {}
+
+  // Base properties for all input types
+  const baseProps = {
+    name,
+    isOptional,
+    ...options
+  }
+
+  // Handle different input types
+  if (isComboInputSpec(inputSpecV1)) {
+    return {
+      type: 'COMBO',
+      ...baseProps,
+      options: isComboInputSpecV1(inputSpecV1)
+        ? inputSpecV1[0]
+        : getComboSpecComboOptions(inputSpecV1)
+    }
+  } else if (typeof inputSpecV1[0] === 'string') {
+    // Handle standard types (INT, FLOAT, BOOLEAN, STRING) and custom types
+    return {
+      type: inputSpecV1[0],
+      ...baseProps
+    }
+  }
+
+  // Fallback for any unhandled cases
+  return {
+    type: 'UNKNOWN',
+    ...baseProps
+  }
+}

src/schemas/nodeDef/nodeDefSchemaV2.ts

@@ -0,0 +1,93 @@
+import { z } from 'zod'
+
+import {
+  zBaseInputOptions,
+  zBooleanInputOptions,
+  zComboInputOptions,
+  zFloatInputOptions,
+  zIntInputOptions,
+  zStringInputOptions
+} from '@/schemas/nodeDefSchema'
+
+const zIntInputSpec = zIntInputOptions.extend({
+  type: z.literal('INT'),
+  name: z.string(),
+  isOptional: z.boolean().optional()
+})
+
+const zFloatInputSpec = zFloatInputOptions.extend({
+  type: z.literal('FLOAT'),
+  name: z.string(),
+  isOptional: z.boolean().optional()
+})
+
+const zBooleanInputSpec = zBooleanInputOptions.extend({
+  type: z.literal('BOOLEAN'),
+  name: z.string(),
+  isOptional: z.boolean().optional()
+})
+
+const zStringInputSpec = zStringInputOptions.extend({
+  type: z.literal('STRING'),
+  name: z.string(),
+  isOptional: z.boolean().optional()
+})
+
+const zComboInputSpec = zComboInputOptions.extend({
+  type: z.literal('COMBO'),
+  name: z.string(),
+  isOptional: z.boolean().optional()
+})
+
+const zCustomInputSpec = zBaseInputOptions.extend({
+  type: z.string(),
+  name: z.string(),
+  isOptional: z.boolean().optional()
+})
+
+const zInputSpec = z.union([
+  zIntInputSpec,
+  zFloatInputSpec,
+  zBooleanInputSpec,
+  zStringInputSpec,
+  zComboInputSpec,
+  zCustomInputSpec
+])
+
+// Output specs
+const zOutputSpec = z.object({
+  index: z.number(),
+  name: z.string(),
+  type: z.string(),
+  is_list: z.boolean(),
+  options: z.array(z.any()).optional(),
+  tooltip: z.string().optional()
+})
+
+// Main node definition schema
+const zNodeDef = z.object({
+  inputs: z.record(zInputSpec),
+  outputs: z.array(zOutputSpec),
+  hidden: z.record(z.any()).optional(),
+
+  name: z.string(),
+  display_name: z.string(),
+  description: z.string(),
+  category: z.string(),
+  output_node: z.boolean(),
+  python_module: z.string(),
+  deprecated: z.boolean().optional(),
+  experimental: z.boolean().optional()
+})
+
+// Export types
+export type IntInputSpec = z.infer<typeof zIntInputSpec>
+export type FloatInputSpec = z.infer<typeof zFloatInputSpec>
+export type BooleanInputSpec = z.infer<typeof zBooleanInputSpec>
+export type StringInputSpec = z.infer<typeof zStringInputSpec>
+export type ComboInputSpec = z.infer<typeof zComboInputSpec>
+export type CustomInputSpec = z.infer<typeof zCustomInputSpec>
+
+export type InputSpec = z.infer<typeof zInputSpec>
+export type OutputSpec = z.infer<typeof zOutputSpec>
+export type ComfyNodeDef = z.infer<typeof zNodeDef>

src/schemas/nodeDefSchema.ts

@@ -13,7 +13,7 @@ const zRemoteWidgetConfig = z.object({
   max_retries: z.number().gte(0).optional()
 })
 
-const zBaseInputOptions = z
+export const zBaseInputOptions = z
   .object({
     default: z.any().optional(),
     /** @deprecated Group node uses this field. Remove when group node feature is removed. */
@@ -28,7 +28,7 @@ const zBaseInputOptions = z
   })
   .passthrough()
 
-const zNumericInputOptions = zBaseInputOptions.extend({
+export const zNumericInputOptions = zBaseInputOptions.extend({
   min: z.number().optional(),
   max: z.number().optional(),
   step: z.number().optional(),
@@ -37,7 +37,7 @@ const zNumericInputOptions = zBaseInputOptions.extend({
   display: z.enum(['slider', 'number', 'knob']).optional()
 })
 
-const zIntInputOptions = zNumericInputOptions.extend({
+export const zIntInputOptions = zNumericInputOptions.extend({
   /**
    * If true, a linked widget will be added to the node to select the mode
    * of `control_after_generate`.
@@ -45,17 +45,17 @@ const zIntInputOptions = zNumericInputOptions.extend({
   control_after_generate: z.boolean().optional()
 })
 
-const zFloatInputOptions = zNumericInputOptions.extend({
+export const zFloatInputOptions = zNumericInputOptions.extend({
   round: z.union([z.number(), z.literal(false)]).optional()
 })
 
-const zBooleanInputOptions = zBaseInputOptions.extend({
+export const zBooleanInputOptions = zBaseInputOptions.extend({
   label_on: z.string().optional(),
   label_off: z.string().optional(),
   default: z.boolean().optional()
 })
 
-const zStringInputOptions = zBaseInputOptions.extend({
+export const zStringInputOptions = zBaseInputOptions.extend({
   default: z.string().optional(),
   multiline: z.boolean().optional(),
   dynamicPrompts: z.boolean().optional(),
@@ -65,7 +65,7 @@ const zStringInputOptions = zBaseInputOptions.extend({
   placeholder: z.string().optional()
 })
 
-const zComboInputOptions = zBaseInputOptions.extend({
+export const zComboInputOptions = zBaseInputOptions.extend({
   control_after_generate: z.boolean().optional(),
   image_upload: z.boolean().optional(),
   image_folder: z.enum(['input', 'output', 'temp']).optional(),

src/services/nodeSearchService.ts

@@ -210,7 +210,7 @@ export class NodeSearchService {
       /* name */ 'Input Type',
       /* invokeSequence */ 'i',
       /* longInvokeSequence */ 'input',
-      (node) => node.inputs.all.map((input) => input.type),
+      (node) => Object.values(node.inputs).map((input) => input.type),
       data,
       filterSearchOptions
     )
@@ -220,7 +220,7 @@ export class NodeSearchService {
       /* name */ 'Output Type',
       /* invokeSequence */ 'o',
       /* longInvokeSequence */ 'output',
-      (node) => node.outputs.all.map((output) => output.type),
+      (node) => node.outputs.map((output) => output.type),
       data,
       filterSearchOptions
     )

src/stores/nodeDefStore.ts

@@ -4,11 +4,16 @@ import { defineStore } from 'pinia'
 import type { TreeNode } from 'primevue/treenode'
 import { computed, ref } from 'vue'
 
-import {
+import { transformNodeDefV1ToV2 } from '@/schemas/nodeDef/migration'
-  type ComfyInputsSpec as ComfyInputsSpecSchema,
+import type {
-  type ComfyNodeDef,
+  ComfyNodeDef as ComfyNodeDefV2,
-  type ComfyOutputTypesSpec as ComfyOutputTypesSpecSchema,
+  InputSpec as InputSpecV2,
-  type InputSpec
+  OutputSpec as OutputSpecV2
+} from '@/schemas/nodeDef/nodeDefSchemaV2'
+import type {
+  ComfyInputsSpec as ComfyInputSpecV1,
+  ComfyNodeDef as ComfyNodeDefV1,
+  ComfyOutputTypesSpec as ComfyOutputSpecV1
 } from '@/schemas/nodeDefSchema'
 import {
   NodeSearchService,
@@ -21,137 +26,8 @@ import {
 } from '@/types/nodeSource'
 import { buildTree } from '@/utils/treeUtil'
 
-export interface BaseInputSpec<T = any> {
+export class ComfyNodeDefImpl implements ComfyNodeDefV1, ComfyNodeDefV2 {
-  name: string
+  // ComfyNodeDef fields (V1)
-  type: string
-  tooltip?: string
-  default?: T
-
-  forceInput?: boolean
-}
-
-export interface NumericInputSpec extends BaseInputSpec<number> {
-  min?: number
-  max?: number
-  step?: number
-}
-
-export interface IntInputSpec extends NumericInputSpec {
-  type: 'INT'
-}
-
-export interface FloatInputSpec extends NumericInputSpec {
-  type: 'FLOAT'
-  round?: number
-}
-
-export interface BooleanInputSpec extends BaseInputSpec<boolean> {
-  type: 'BOOLEAN'
-  labelOn?: string
-  labelOff?: string
-}
-
-export interface StringInputSpec extends BaseInputSpec<string> {
-  type: 'STRING'
-  multiline?: boolean
-  dynamicPrompts?: boolean
-}
-
-export interface ComboInputSpec extends BaseInputSpec<any> {
-  type: 'COMBO'
-  comboOptions: any[]
-  controlAfterGenerate?: boolean
-  imageUpload?: boolean
-}
-
-export class ComfyInputsSpec {
-  required: Record<string, BaseInputSpec>
-  optional: Record<string, BaseInputSpec>
-  hidden?: Record<string, any>
-
-  constructor(obj: ComfyInputsSpecSchema) {
-    this.required = ComfyInputsSpec.transformInputSpecRecord(obj.required ?? {})
-    this.optional = ComfyInputsSpec.transformInputSpecRecord(obj.optional ?? {})
-    this.hidden = obj.hidden
-  }
-
-  private static transformInputSpecRecord(
-    record: Record<string, InputSpec>
-  ): Record<string, BaseInputSpec> {
-    const result: Record<string, BaseInputSpec> = {}
-    for (const [key, value] of Object.entries(record)) {
-      result[key] = ComfyInputsSpec.transformSingleInputSpec(key, value)
-    }
-    return result
-  }
-
-  private static isInputSpec(obj: any): boolean {
-    return (
-      Array.isArray(obj) &&
-      obj.length >= 1 &&
-      (typeof obj[0] === 'string' || Array.isArray(obj[0]))
-    )
-  }
-
-  private static transformSingleInputSpec(
-    name: string,
-    value: any
-  ): BaseInputSpec {
-    if (!ComfyInputsSpec.isInputSpec(value)) return value
-
-    const [typeRaw, _spec] = value
-    const spec = _spec ?? {}
-    const type = Array.isArray(typeRaw) ? 'COMBO' : value[0]
-
-    switch (type) {
-      case 'COMBO':
-        return {
-          name,
-          type,
-          ...spec,
-          comboOptions: typeRaw,
-          default: spec.default ?? typeRaw[0]
-        } as ComboInputSpec
-      case 'INT':
-      case 'FLOAT':
-      case 'BOOLEAN':
-      case 'STRING':
-      default:
-        return { name, type, ...spec } as BaseInputSpec
-    }
-  }
-
-  get all() {
-    return [...Object.values(this.required), ...Object.values(this.optional)]
-  }
-
-  getInput(name: string): BaseInputSpec | undefined {
-    return this.required[name] ?? this.optional[name]
-  }
-}
-
-export class ComfyOutputSpec {
-  constructor(
-    public index: number,
-    // Name is not unique for output params
-    public name: string,
-    public type: string,
-    public is_list: boolean,
-    public comboOptions?: any[],
-    public tooltip?: string
-  ) {}
-}
-
-export class ComfyOutputsSpec {
-  constructor(public outputs: ComfyOutputSpec[]) {}
-
-  get all() {
-    return this.outputs
-  }
-}
-
-export class ComfyNodeDefImpl implements ComfyNodeDef {
-  // ComfyNodeDef fields
   readonly name: string
   readonly display_name: string
   /**
@@ -167,11 +43,11 @@ export class ComfyNodeDefImpl implements ComfyNodeDef {
   /**
    * @deprecated Use `inputs` instead
    */
-  readonly input: ComfyInputsSpecSchema
+  readonly input: ComfyInputSpecV1
   /**
    * @deprecated Use `outputs` instead
    */
-  readonly output: ComfyOutputTypesSpecSchema
+  readonly output: ComfyOutputSpecV1
   /**
    * @deprecated Use `outputs[n].is_list` instead
    */
@@ -185,12 +61,16 @@ export class ComfyNodeDefImpl implements ComfyNodeDef {
    */
   readonly output_tooltips?: string[]
 
+  // V2 fields
+  readonly inputs: Record<string, InputSpecV2>
+  readonly outputs: OutputSpecV2[]
+  readonly hidden?: Record<string, any>
+
   // ComfyNodeDefImpl fields
-  readonly inputs: ComfyInputsSpec
-  readonly outputs: ComfyOutputsSpec
   readonly nodeSource: NodeSource
 
-  constructor(obj: ComfyNodeDef) {
+  constructor(obj: ComfyNodeDefV1) {
+    // Initialize V1 fields
     this.name = obj.name
     this.display_name = obj.display_name
     this.category = obj.category
@@ -206,28 +86,16 @@ export class ComfyNodeDefImpl implements ComfyNodeDef {
     this.output_name = obj.output_name
     this.output_tooltips = obj.output_tooltips
 
-    this.inputs = new ComfyInputsSpec(obj.input ?? {})
+    // Initialize V2 fields
-    this.outputs = ComfyNodeDefImpl.transformOutputSpec(obj)
+    const defV2 = transformNodeDefV1ToV2(obj)
+    this.inputs = defV2.inputs
+    this.outputs = defV2.outputs
+    this.hidden = defV2.hidden
+
+    // Initialize node source
     this.nodeSource = getNodeSource(obj.python_module)
   }
 
-  private static transformOutputSpec(obj: any): ComfyOutputsSpec {
-    const { output, output_is_list, output_name, output_tooltips } = obj
-    const result = (output ?? []).map((type: string | any[], index: number) => {
-      const typeString = Array.isArray(type) ? 'COMBO' : type
-
-      return new ComfyOutputSpec(
-        index,
-        output_name?.[index],
-        typeString,
-        output_is_list?.[index],
-        Array.isArray(type) ? type : undefined,
-        output_tooltips?.[index]
-      )
-    })
-    return new ComfyOutputsSpec(result)
-  }
-
   get nodePath(): string {
     return (this.category ? this.category + '/' : '') + this.name
   }
@@ -253,7 +121,7 @@ export class ComfyNodeDefImpl implements ComfyNodeDef {
   }
 }
 
-export const SYSTEM_NODE_DEFS: Record<string, ComfyNodeDef> = {
+export const SYSTEM_NODE_DEFS: Record<string, ComfyNodeDefV1> = {
   PrimitiveNode: {
     name: 'PrimitiveNode',
     display_name: 'Primitive',
@@ -323,7 +191,7 @@ export function createDummyFolderNodeDef(folderPath: string): ComfyNodeDefImpl {
     output_name: [],
     output_is_list: [],
     output_node: false
-  } as ComfyNodeDef)
+  } as ComfyNodeDefV1)
 }
 
 export const useNodeDefStore = defineStore('nodeDef', () => {
@@ -336,10 +204,10 @@ export const useNodeDefStore = defineStore('nodeDef', () => {
   const nodeDataTypes = computed(() => {
     const types = new Set<string>()
     for (const nodeDef of nodeDefs.value) {
-      for (const input of nodeDef.inputs.all) {
+      for (const input of Object.values(nodeDef.inputs)) {
         types.add(input.type)
       }
-      for (const output of nodeDef.outputs.all) {
+      for (const output of nodeDef.outputs) {
         types.add(output.type)
       }
     }
@@ -357,7 +225,7 @@ export const useNodeDefStore = defineStore('nodeDef', () => {
   )
   const nodeTree = computed(() => buildNodeDefTree(visibleNodeDefs.value))
 
-  function updateNodeDefs(nodeDefs: ComfyNodeDef[]) {
+  function updateNodeDefs(nodeDefs: ComfyNodeDefV1[]) {
     const newNodeDefsByName: Record<string, ComfyNodeDefImpl> = {}
     const newNodeDefsByDisplayName: Record<string, ComfyNodeDefImpl> = {}
     for (const nodeDef of nodeDefs) {
@@ -374,7 +242,7 @@ export const useNodeDefStore = defineStore('nodeDef', () => {
     nodeDefsByName.value = newNodeDefsByName
     nodeDefsByDisplayName.value = newNodeDefsByDisplayName
   }
-  function addNodeDef(nodeDef: ComfyNodeDef) {
+  function addNodeDef(nodeDef: ComfyNodeDefV1) {
     const nodeDefImpl = new ComfyNodeDefImpl(nodeDef)
     nodeDefsByName.value[nodeDef.name] = nodeDefImpl
     nodeDefsByDisplayName.value[nodeDef.display_name] = nodeDefImpl

src/stores/widgetStore.ts

@@ -1,6 +1,7 @@
 import { defineStore } from 'pinia'
 import { computed, ref } from 'vue'
 
+import type { InputSpec as InputSpecV2 } from '@/schemas/nodeDef/nodeDefSchemaV2'
 import {
   type ComboInputSpecV2,
   type InputSpec,
@@ -8,8 +9,6 @@ import {
 } from '@/schemas/nodeDefSchema'
 import { ComfyWidgetConstructor, ComfyWidgets } from '@/scripts/widgets'
 
-import type { BaseInputSpec } from './nodeDefStore'
-
 export const useWidgetStore = defineStore('widget', () => {
   const coreWidgets = ComfyWidgets
   const customWidgets = ref<Record<string, ComfyWidgetConstructor>>({})
@@ -33,7 +32,7 @@ export const useWidgetStore = defineStore('widget', () => {
     }
   }
 
-  function inputIsWidget(spec: BaseInputSpec) {
+  function inputIsWidget(spec: InputSpecV2) {
     return getWidgetType(spec.type, spec.name) !== null
   }
 

tests-ui/tests/nodeDef.test.ts

@@ -1,18 +1,12 @@
 // @ts-strict-ignore
 import { describe, expect, it } from 'vitest'
 
-import type { ComfyNodeDef } from '@/schemas/nodeDefSchema'
+import { transformNodeDefV1ToV2 } from '@/schemas/nodeDef/migration'
-import {
+import type { ComfyNodeDef as ComfyNodeDefV1 } from '@/schemas/nodeDefSchema'
-  BooleanInputSpec,
+import { ComfyNodeDefImpl } from '@/stores/nodeDefStore'
-  ComfyInputsSpec,
-  ComfyNodeDefImpl,
-  FloatInputSpec,
-  IntInputSpec,
-  StringInputSpec
-} from '@/stores/nodeDefStore'
 
-describe('ComfyInputsSpec', () => {
+describe('NodeDef Migration', () => {
-  it('should transform a plain object to ComfyInputsSpec instance', () => {
+  it('should transform a plain object to V2 format', () => {
     const plainObject = {
       required: {
         intInput: ['INT', { min: 0, max: 100, default: 50 }],
@@ -28,13 +22,26 @@ describe('ComfyInputsSpec', () => {
       hidden: {
         someHiddenValue: 42
       }
-    } as ComfyNodeDef['input']
+    } as ComfyNodeDefV1['input']
 
-    const result = new ComfyInputsSpec(plainObject)
+    const nodeDef: ComfyNodeDefV1 = {
+      name: 'TestNode',
+      display_name: 'Test Node',
+      category: 'Testing',
+      python_module: 'test_module',
+      description: 'A test node',
+      input: plainObject,
+      output: ['INT'],
+      output_is_list: [false],
+      output_name: ['intOutput'],
+      output_node: false
+    }
 
-    expect(result).toBeInstanceOf(ComfyInputsSpec)
+    const result = transformNodeDefV1ToV2(nodeDef)
-    expect(result.required).toBeDefined()
+
-    expect(result.optional).toBeDefined()
+    expect(result).toBeDefined()
+    expect(result.inputs).toBeDefined()
+    expect(result.outputs).toBeDefined()
     expect(result.hidden).toBeDefined()
   })
 
@@ -44,20 +51,35 @@ describe('ComfyInputsSpec', () => {
         intInput: ['INT', { min: 0, max: 100, default: 50 }],
         stringInput: ['STRING', { default: 'Hello', multiline: true }]
       }
-    } as ComfyNodeDef['input']
+    } as ComfyNodeDefV1['input']
 
-    const result = new ComfyInputsSpec(plainObject)
+    const nodeDef: ComfyNodeDefV1 = {
+      name: 'TestNode',
+      display_name: 'Test Node',
+      category: 'Testing',
+      python_module: 'test_module',
+      description: 'A test node',
+      input: plainObject,
+      output: [],
+      output_is_list: [],
+      output_name: [],
+      output_node: false
+    }
 
-    const intInput = result.required.intInput as IntInputSpec
+    const result = transformNodeDefV1ToV2(nodeDef)
-    const stringInput = result.required.stringInput as StringInputSpec
+
+    const intInput = result.inputs['intInput']
+    const stringInput = result.inputs['stringInput']
 
     expect(intInput.min).toBe(0)
     expect(intInput.max).toBe(100)
     expect(intInput.default).toBe(50)
     expect(intInput.name).toBe('intInput')
+    expect(intInput.type).toBe('INT')
     expect(stringInput.default).toBe('Hello')
     expect(stringInput.multiline).toBe(true)
     expect(stringInput.name).toBe('stringInput')
+    expect(stringInput.type).toBe('STRING')
   })
 
   it('should correctly transform optional input specs', () => {
@@ -69,19 +91,36 @@ describe('ComfyInputsSpec', () => {
         ],
         floatInput: ['FLOAT', { min: 0, max: 1, step: 0.1 }]
       }
-    } as ComfyNodeDef['input']
+    } as ComfyNodeDefV1['input']
 
-    const result = new ComfyInputsSpec(plainObject)
+    const nodeDef: ComfyNodeDefV1 = {
+      name: 'TestNode',
+      display_name: 'Test Node',
+      category: 'Testing',
+      python_module: 'test_module',
+      description: 'A test node',
+      input: plainObject,
+      output: [],
+      output_is_list: [],
+      output_name: [],
+      output_node: false
+    }
 
-    const booleanInput = result.optional.booleanInput as BooleanInputSpec
+    const result = transformNodeDefV1ToV2(nodeDef)
-    const floatInput = result.optional.floatInput as FloatInputSpec
+
+    const booleanInput = result.inputs['booleanInput']
+    const floatInput = result.inputs['floatInput']
 
     expect(booleanInput.default).toBe(true)
     expect(booleanInput.labelOn).toBe('Yes')
     expect(booleanInput.labelOff).toBe('No')
+    expect(booleanInput.type).toBe('BOOLEAN')
+    expect(booleanInput.isOptional).toBe(true)
     expect(floatInput.min).toBe(0)
     expect(floatInput.max).toBe(1)
     expect(floatInput.step).toBe(0.1)
+    expect(floatInput.type).toBe('FLOAT')
+    expect(floatInput.isOptional).toBe(true)
   })
 
   it('should handle combo input specs', () => {
@@ -89,11 +128,25 @@ describe('ComfyInputsSpec', () => {
       optional: {
         comboInput: [[1, 2, 3], { default: 2 }]
       }
-    } as ComfyNodeDef['input']
+    } as ComfyNodeDefV1['input']
 
-    const result = new ComfyInputsSpec(plainObject)
+    const nodeDef: ComfyNodeDefV1 = {
-    expect(result.optional.comboInput.type).toBe('COMBO')
+      name: 'TestNode',
-    expect(result.optional.comboInput.default).toBe(2)
+      display_name: 'Test Node',
+      category: 'Testing',
+      python_module: 'test_module',
+      description: 'A test node',
+      input: plainObject,
+      output: [],
+      output_is_list: [],
+      output_name: [],
+      output_node: false
+    }
+
+    const result = transformNodeDefV1ToV2(nodeDef)
+    expect(result.inputs['comboInput'].type).toBe('COMBO')
+    expect(result.inputs['comboInput'].default).toBe(2)
+    expect(result.inputs['comboInput'].options).toEqual([1, 2, 3])
   })
 
   it('should handle combo input specs (auto-default)', () => {
@@ -101,12 +154,25 @@ describe('ComfyInputsSpec', () => {
       optional: {
         comboInput: [[1, 2, 3], {}]
       }
-    } as ComfyNodeDef['input']
+    } as ComfyNodeDefV1['input']
 
-    const result = new ComfyInputsSpec(plainObject)
+    const nodeDef: ComfyNodeDefV1 = {
-    expect(result.optional.comboInput.type).toBe('COMBO')
+      name: 'TestNode',
+      display_name: 'Test Node',
+      category: 'Testing',
+      python_module: 'test_module',
+      description: 'A test node',
+      input: plainObject,
+      output: [],
+      output_is_list: [],
+      output_name: [],
+      output_node: false
+    }
+
+    const result = transformNodeDefV1ToV2(nodeDef)
+    expect(result.inputs['comboInput'].type).toBe('COMBO')
     // Should pick the first choice as default
-    expect(result.optional.comboInput.default).toBe(1)
+    expect(result.inputs['comboInput'].options).toEqual([1, 2, 3])
   })
 
   it('should handle custom input specs', () => {
@@ -114,11 +180,24 @@ describe('ComfyInputsSpec', () => {
       optional: {
         customInput: ['CUSTOM_TYPE', { default: 'custom value' }]
       }
-    } as ComfyNodeDef['input']
+    } as ComfyNodeDefV1['input']
 
-    const result = new ComfyInputsSpec(plainObject)
+    const nodeDef: ComfyNodeDefV1 = {
-    expect(result.optional.customInput.type).toBe('CUSTOM_TYPE')
+      name: 'TestNode',
-    expect(result.optional.customInput.default).toBe('custom value')
+      display_name: 'Test Node',
+      category: 'Testing',
+      python_module: 'test_module',
+      description: 'A test node',
+      input: plainObject,
+      output: [],
+      output_is_list: [],
+      output_name: [],
+      output_node: false
+    }
+
+    const result = transformNodeDefV1ToV2(nodeDef)
+    expect(result.inputs['customInput'].type).toBe('CUSTOM_TYPE')
+    expect(result.inputs['customInput'].default).toBe('custom value')
   })
 
   it('should not transform hidden fields', () => {
@@ -127,9 +206,22 @@ describe('ComfyInputsSpec', () => {
         someHiddenValue: 42,
         anotherHiddenValue: { nested: 'object' }
       }
-    } as ComfyNodeDef['input']
+    } as ComfyNodeDefV1['input']
 
-    const result = new ComfyInputsSpec(plainObject)
+    const nodeDef: ComfyNodeDefV1 = {
+      name: 'TestNode',
+      display_name: 'Test Node',
+      category: 'Testing',
+      python_module: 'test_module',
+      description: 'A test node',
+      input: plainObject,
+      output: [],
+      output_is_list: [],
+      output_name: [],
+      output_node: false
+    }
+
+    const result = transformNodeDefV1ToV2(nodeDef)
 
     expect(result.hidden).toEqual(plainObject.hidden)
     expect(result.hidden?.someHiddenValue).toBe(42)
@@ -139,11 +231,23 @@ describe('ComfyInputsSpec', () => {
   it('should handle empty or undefined fields', () => {
     const plainObject = {}
 
-    const result = new ComfyInputsSpec(plainObject)
+    const nodeDef: ComfyNodeDefV1 = {
+      name: 'TestNode',
+      display_name: 'Test Node',
+      category: 'Testing',
+      python_module: 'test_module',
+      description: 'A test node',
+      input: plainObject,
+      output: [],
+      output_is_list: [],
+      output_name: [],
+      output_node: false
+    }
 
-    expect(result).toBeInstanceOf(ComfyInputsSpec)
+    const result = transformNodeDefV1ToV2(nodeDef)
-    expect(result.required).toEqual({})
+
-    expect(result.optional).toEqual({})
+    expect(result).toBeDefined()
+    expect(result.inputs).toEqual({})
     expect(result.hidden).toBeUndefined()
   })
 })
@@ -163,8 +267,9 @@ describe('ComfyNodeDefImpl', () => {
       },
       output: ['INT'],
       output_is_list: [false],
-      output_name: ['intOutput']
+      output_name: ['intOutput'],
-    } as ComfyNodeDef
+      output_node: false
+    } as ComfyNodeDefV1
 
     const result = new ComfyNodeDefImpl(plainObject)
 
@@ -174,8 +279,8 @@ describe('ComfyNodeDefImpl', () => {
     expect(result.category).toBe('Testing')
     expect(result.python_module).toBe('test_module')
     expect(result.description).toBe('A test node')
-    expect(result.inputs).toBeInstanceOf(ComfyInputsSpec)
+    expect(result.inputs).toBeDefined()
-    expect(result.outputs.all).toEqual([
+    expect(result.outputs).toEqual([
       {
         index: 0,
         name: 'intOutput',
@@ -201,8 +306,9 @@ describe('ComfyNodeDefImpl', () => {
       output: ['INT'],
       output_is_list: [false],
       output_name: ['intOutput'],
-      deprecated: true
+      deprecated: true,
-    } as ComfyNodeDef
+      output_node: false
+    } as ComfyNodeDefV1
 
     const result = new ComfyNodeDefImpl(plainObject)
     expect(result.deprecated).toBe(true)
@@ -224,8 +330,9 @@ describe('ComfyNodeDefImpl', () => {
       },
       output: ['INT'],
       output_is_list: [false],
-      output_name: ['intOutput']
+      output_name: ['intOutput'],
-    } as ComfyNodeDef
+      output_node: false
+    } as ComfyNodeDefV1
 
     const result = new ComfyNodeDefImpl(plainObject)
     expect(result.deprecated).toBe(true)
@@ -241,12 +348,13 @@ describe('ComfyNodeDefImpl', () => {
       input: {},
       output: ['STRING', ['COMBO', 'option1', 'option2'], 'FLOAT'],
       output_is_list: [true, false, false],
-      output_name: ['stringOutput', 'comboOutput', 'floatOutput']
+      output_name: ['stringOutput', 'comboOutput', 'floatOutput'],
-    }
+      output_node: false
+    } as ComfyNodeDefV1
 
     const result = new ComfyNodeDefImpl(plainObject)
 
-    expect(result.outputs.all).toEqual([
+    expect(result.outputs).toEqual([
       {
         index: 0,
         name: 'stringOutput',
@@ -258,7 +366,7 @@ describe('ComfyNodeDefImpl', () => {
         name: 'comboOutput',
         type: 'COMBO',
         is_list: false,
-        comboOptions: ['COMBO', 'option1', 'option2']
+        options: ['COMBO', 'option1', 'option2']
       },
       {
         index: 2,
@@ -279,12 +387,13 @@ describe('ComfyNodeDefImpl', () => {
       input: {},
       output: ['INT', 'FLOAT', 'FLOAT'],
       output_is_list: [false, true, true],
-      output_name: ['INT', 'FLOAT', 'FLOAT']
+      output_name: ['INT', 'FLOAT', 'FLOAT'],
-    }
+      output_node: false
+    } as ComfyNodeDefV1
 
     const result = new ComfyNodeDefImpl(plainObject)
 
-    expect(result.outputs.all).toEqual([
+    expect(result.outputs).toEqual([
       {
         index: 0,
         name: 'INT',
@@ -316,11 +425,12 @@ describe('ComfyNodeDefImpl', () => {
       input: {},
       output: ['INT', 'FLOAT', 'STRING'],
       output_is_list: [false, false, false],
-      output_name: ['output', 'output', 'uniqueOutput']
+      output_name: ['output', 'output', 'uniqueOutput'],
-    }
+      output_node: false
+    } as ComfyNodeDefV1
 
     const result = new ComfyNodeDefImpl(plainObject)
-    expect(result.outputs.all).toEqual([
+    expect(result.outputs).toEqual([
       {
         index: 0,
         name: 'output',
@@ -352,12 +462,13 @@ describe('ComfyNodeDefImpl', () => {
       input: {},
       output: [],
       output_is_list: [],
-      output_name: []
+      output_name: [],
-    }
+      output_node: false
+    } as ComfyNodeDefV1
 
     const result = new ComfyNodeDefImpl(plainObject)
 
-    expect(result.outputs.all).toEqual([])
+    expect(result.outputs).toEqual([])
   })
 
   it('should handle undefined fields', () => {
@@ -366,12 +477,13 @@ describe('ComfyNodeDefImpl', () => {
       display_name: 'Empty Output Node',
       category: 'Test',
       python_module: 'test_module',
-      description: 'A node with no outputs'
+      description: 'A node with no outputs',
-    }
+      output_node: false
+    } as ComfyNodeDefV1
 
     const result = new ComfyNodeDefImpl(plainObject)
-    expect(result.outputs.all).toEqual([])
+    expect(result.outputs).toEqual([])
-    expect(result.inputs.all).toEqual([])
+    expect(Object.keys(result.inputs)).toHaveLength(0)
   })
 
   it('should handle complex input specifications', () => {
@@ -393,13 +505,17 @@ describe('ComfyNodeDefImpl', () => {
       },
       output: ['INT'],
       output_is_list: [false],
-      output_name: ['result']
+      output_name: ['result'],
-    } as ComfyNodeDef
+      output_node: false
+    } as ComfyNodeDefV1
 
     const result = new ComfyNodeDefImpl(plainObject)
 
-    expect(result.inputs).toBeInstanceOf(ComfyInputsSpec)
+    expect(result.inputs).toBeDefined()
-    expect(result.inputs.required).toBeDefined()
+    expect(Object.keys(result.inputs)).toHaveLength(4)
-    expect(result.inputs.optional).toBeDefined()
+    expect(result.inputs['intInput']).toBeDefined()
+    expect(result.inputs['stringInput']).toBeDefined()
+    expect(result.inputs['booleanInput']).toBeDefined()
+    expect(result.inputs['floatInput']).toBeDefined()
   })
 })