[Refactor] Use NodeSlot.hasErrors API (#3284)

Co-authored-by: github-actions <github-actions@github.com>

src/components/graph/GraphCanvas.vue

@@ -37,6 +37,7 @@
 </template>
 
 <script setup lang="ts">
+import type { LGraphNode } from '@comfyorg/litegraph'
 import { computed, onMounted, ref, watch, watchEffect } from 'vue'
 
 import LiteGraphCanvasSplitterOverlay from '@/components/LiteGraphCanvasSplitterOverlay.vue'
@@ -179,6 +180,37 @@ watch(
   }
 )
 
+// Update node slot errors
+watch(
+  () => executionStore.lastNodeErrors,
+  (lastNodeErrors) => {
+    const removeSlotError = (node: LGraphNode) => {
+      for (const slot of node.inputs) {
+        delete slot.hasErrors
+      }
+      for (const slot of node.outputs) {
+        delete slot.hasErrors
+      }
+    }
+
+    for (const node of comfyApp.graph.nodes) {
+      removeSlotError(node)
+      const nodeErrors = lastNodeErrors?.[node.id]
+      if (!nodeErrors) continue
+      for (const error of nodeErrors.errors) {
+        if (error.extra_info && error.extra_info.input_name) {
+          const inputIndex = node.findInputSlot(error.extra_info.input_name)
+          if (inputIndex !== -1) {
+            node.inputs[inputIndex].hasErrors = true
+          }
+        }
+      }
+    }
+
+    comfyApp.canvas.draw(true, true)
+  }
+)
+
 const loadCustomNodesI18n = async () => {
   try {
     const i18nData = await api.getCustomNodesI18n()

src/schemas/apiSchema.ts

@@ -259,7 +259,12 @@ const zError = z.object({
   type: z.string(),
   message: z.string(),
   details: z.string(),
-  extra_info: z.record(z.string(), z.any())
+  extra_info: z
+    .object({
+      input_name: z.string().optional()
+    })
+    .passthrough()
+    .optional()
 })
 const zNodeError = z.object({
   errors: z.array(zError),

src/scripts/app.ts

@@ -11,7 +11,11 @@ import type { ToastMessageOptions } from 'primevue/toast'
 import { reactive } from 'vue'
 
 import { st, t } from '@/i18n'
-import type { NodeError, ResultItem } from '@/schemas/apiSchema'
+import type {
+  ExecutionErrorWsMessage,
+  NodeError,
+  ResultItem
+} from '@/schemas/apiSchema'
 import {
   ComfyApiWorkflow,
   type ComfyWorkflowJSON,
@@ -120,9 +124,6 @@ export class ComfyApp {
   dragOverNode: LGraphNode | null = null
   // @ts-expect-error fixme ts strict error
   canvasEl: HTMLCanvasElement
-  lastNodeErrors: Record<NodeId, NodeError> | null = null
-  /** @type {ExecutionErrorWsMessage} */
-  lastExecutionError: { node_id?: NodeId } | null = null
   configuringGraph: boolean = false
   // @ts-expect-error fixme ts strict error
   ctx: CanvasRenderingContext2D
@@ -136,6 +137,22 @@ export class ComfyApp {
   // Set by Comfy.Clipspace extension
   openClipspace: () => void = () => {}
 
+  /**
+   * The node errors from the previous execution.
+   * @deprecated Use useExecutionStore().lastNodeErrors instead
+   */
+  get lastNodeErrors(): Record<NodeId, NodeError> | null {
+    return useExecutionStore().lastNodeErrors
+  }
+
+  /**
+   * The error from the previous execution.
+   * @deprecated Use useExecutionStore().lastExecutionError instead
+   */
+  get lastExecutionError(): ExecutionErrorWsMessage | null {
+    return useExecutionStore().lastExecutionError
+  }
+
   /**
    * @deprecated Use useExecutionStore().executingNodeId instead
    */
@@ -553,51 +570,7 @@ export class ComfyApp {
     }
   }
 
-  /**
-   * Draws node highlights (executing, drag drop) and progress bar
-   */
   #addDrawNodeHandler() {
-    const origDrawNodeShape = LGraphCanvas.prototype.drawNodeShape
-    const self = this
-    LGraphCanvas.prototype.drawNodeShape = function (
-      node,
-      ctx,
-      _size,
-      _fgcolor,
-      _bgcolor
-    ) {
-      // @ts-expect-error fixme ts strict error
-      const res = origDrawNodeShape.apply(this, arguments)
-
-      const nodeErrors = self.lastNodeErrors?.[node.id]
-
-      // Highlight inputs that failed validation
-      if (nodeErrors) {
-        ctx.lineWidth = 2
-        ctx.strokeStyle = 'red'
-        for (const error of nodeErrors.errors) {
-          if (error.extra_info && error.extra_info.input_name) {
-            const inputIndex = node.findInputSlot(error.extra_info.input_name)
-            if (inputIndex !== -1) {
-              let pos = node.getConnectionPos(true, inputIndex)
-              ctx.beginPath()
-              ctx.arc(
-                pos[0] - node.pos[0],
-                pos[1] - node.pos[1],
-                12,
-                0,
-                2 * Math.PI,
-                false
-              )
-              ctx.stroke()
-            }
-          }
-        }
-      }
-
-      return res
-    }
-
     const origDrawNode = LGraphCanvas.prototype.drawNode
     LGraphCanvas.prototype.drawNode = function (node) {
       const editor_alpha = this.editor_alpha
@@ -684,14 +657,12 @@ export class ComfyApp {
     })
 
     api.addEventListener('execution_start', () => {
-      this.lastExecutionError = null
       this.graph.nodes.forEach((node) => {
         if (node.onExecutionStart) node.onExecutionStart()
       })
     })
 
     api.addEventListener('execution_error', ({ detail }) => {
-      this.lastExecutionError = detail
       useDialogService().showExecutionErrorDialog(detail)
       this.canvas.draw(true, true)
     })
@@ -1183,7 +1154,8 @@ export class ComfyApp {
     }
 
     this.#processingQueue = true
-    this.lastNodeErrors = null
+    const executionStore = useExecutionStore()
+    executionStore.lastNodeErrors = null
 
     try {
       while (this.#queueItems.length) {
@@ -1197,13 +1169,13 @@ export class ComfyApp {
           const p = await this.graphToPrompt()
           try {
             const res = await api.queuePrompt(number, p)
-            this.lastNodeErrors = res.node_errors ?? null
-            if (this.lastNodeErrors?.length) {
+            executionStore.lastNodeErrors = res.node_errors ?? null
+            if (executionStore.lastNodeErrors?.length) {
               this.canvas.draw(true, true)
             } else {
               try {
                 if (res.prompt_id) {
-                  useExecutionStore().storePrompt({
+                  executionStore.storePrompt({
                     id: res.prompt_id,
                     nodes: Object.keys(p.output),
                     workflow: useWorkspaceStore().workflow
@@ -1219,7 +1191,7 @@ export class ComfyApp {
             })
 
             if (error instanceof PromptExecutionError) {
-              this.lastNodeErrors = error.response.node_errors ?? null
+              executionStore.lastNodeErrors = error.response.node_errors ?? null
               this.canvas.draw(true, true)
             }
             break
@@ -1241,7 +1213,7 @@ export class ComfyApp {
       this.#processingQueue = false
     }
     api.dispatchCustomEvent('promptQueued', { number, batchCount })
-    return !this.lastNodeErrors
+    return !executionStore.lastNodeErrors
   }
 
   showErrorOnFileLoad(file: File) {
@@ -1569,8 +1541,9 @@ export class ComfyApp {
       this.revokePreviews(id)
     }
     this.nodePreviewImages = {}
-    this.lastNodeErrors = null
-    this.lastExecutionError = null
+    const executionStore = useExecutionStore()
+    executionStore.lastNodeErrors = null
+    executionStore.lastExecutionError = null
   }
 
   clientPosToCanvasPos(pos: Vector2): Vector2 {

src/stores/executionStore.ts

@@ -5,7 +5,9 @@ import type {
   ExecutedWsMessage,
   ExecutingWsMessage,
   ExecutionCachedWsMessage,
+  ExecutionErrorWsMessage,
   ExecutionStartWsMessage,
+  NodeError,
   ProgressWsMessage
 } from '@/schemas/apiSchema'
 import type {
@@ -33,6 +35,8 @@ export const useExecutionStore = defineStore('execution', () => {
   const clientId = ref<string | null>(null)
   const activePromptId = ref<string | null>(null)
   const queuedPrompts = ref<Record<NodeId, QueuedPrompt>>({})
+  const lastNodeErrors = ref<Record<NodeId, NodeError> | null>(null)
+  const lastExecutionError = ref<ExecutionErrorWsMessage | null>(null)
   const executingNodeId = ref<string | null>(null)
   const executingNode = computed<ComfyNode | null>(() => {
     if (!executingNodeId.value) return null
@@ -95,6 +99,10 @@ export const useExecutionStore = defineStore('execution', () => {
     api.addEventListener('executing', handleExecuting as EventListener)
     api.addEventListener('progress', handleProgress as EventListener)
     api.addEventListener('status', handleStatus as EventListener)
+    api.addEventListener(
+      'execution_error',
+      handleExecutionError as EventListener
+    )
   }
 
   function unbindExecutionEvents() {
@@ -110,9 +118,14 @@ export const useExecutionStore = defineStore('execution', () => {
     api.removeEventListener('executing', handleExecuting as EventListener)
     api.removeEventListener('progress', handleProgress as EventListener)
     api.removeEventListener('status', handleStatus as EventListener)
+    api.removeEventListener(
+      'execution_error',
+      handleExecutionError as EventListener
+    )
   }
 
   function handleExecutionStart(e: CustomEvent<ExecutionStartWsMessage>) {
+    lastExecutionError.value = null
     activePromptId.value = e.detail.prompt_id
     queuedPrompts.value[activePromptId.value] ??= { nodes: {} }
   }
@@ -161,6 +174,10 @@ export const useExecutionStore = defineStore('execution', () => {
     }
   }
 
+  function handleExecutionError(e: CustomEvent<ExecutionErrorWsMessage>) {
+    lastExecutionError.value = e.detail
+  }
+
   function storePrompt({
     nodes,
     id,
@@ -197,6 +214,14 @@ export const useExecutionStore = defineStore('execution', () => {
      * The queued prompts
      */
     queuedPrompts,
+    /**
+     * The node errors from the previous execution.
+     */
+    lastNodeErrors,
+    /**
+     * The error from the previous execution.
+     */
+    lastExecutionError,
     /**
      * The id of the node that is currently being executed
      */