From 8c789bd05d962463a6b4e3a66daa13ae7b1f4087 Mon Sep 17 00:00:00 2001
From: Chenlei Hu <huchenlei@proton.me>
Date: Sun, 30 Mar 2025 12:05:45 -0400
Subject: [PATCH] [Refactor] Use litegraph LGraphNode.strokeStyles API (#3278)

---
 package-lock.json                |  8 +++----
 package.json                     |  2 +-
 src/scripts/app.ts               | 40 ++------------------------------
 src/services/litegraphService.ts | 20 ++++++++++++++++
 4 files changed, 27 insertions(+), 43 deletions(-)

diff --git a/package-lock.json b/package-lock.json
index 04ebc879d..8de29408e 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -12,7 +12,7 @@
         "@alloc/quick-lru": "^5.2.0",
         "@atlaskit/pragmatic-drag-and-drop": "^1.3.1",
         "@comfyorg/comfyui-electron-types": "^0.4.31",
-        "@comfyorg/litegraph": "^0.11.5",
+        "@comfyorg/litegraph": "^0.11.6",
         "@primevue/forms": "^4.2.5",
         "@primevue/themes": "^4.2.5",
         "@sentry/vue": "^8.48.0",
@@ -478,9 +478,9 @@
       "license": "GPL-3.0-only"
     },
     "node_modules/@comfyorg/litegraph": {
-      "version": "0.11.5",
-      "resolved": "https://registry.npmjs.org/@comfyorg/litegraph/-/litegraph-0.11.5.tgz",
-      "integrity": "sha512-GBNloRNLEI/R+7zFcrXauMND1VJAMYqjheX6hcbqucCzSHBNlkcJnbhstvILSPIuv0GP1tEgzN4EDe5JtKIqoQ==",
+      "version": "0.11.6",
+      "resolved": "https://registry.npmjs.org/@comfyorg/litegraph/-/litegraph-0.11.6.tgz",
+      "integrity": "sha512-X2VuqX8ipX/49WeaQpJ03NTb0lHR91KG2bPaMokqRh+64D9IIwe703dmuBSIxvF4ZfFMLCchcsjsRJcC6Bj/5w==",
       "license": "MIT"
     },
     "node_modules/@cspotcode/source-map-support": {
diff --git a/package.json b/package.json
index 3a5cdc57b..708d528b8 100644
--- a/package.json
+++ b/package.json
@@ -71,7 +71,7 @@
     "@alloc/quick-lru": "^5.2.0",
     "@atlaskit/pragmatic-drag-and-drop": "^1.3.1",
     "@comfyorg/comfyui-electron-types": "^0.4.31",
-    "@comfyorg/litegraph": "^0.11.5",
+    "@comfyorg/litegraph": "^0.11.6",
     "@primevue/forms": "^4.2.5",
     "@primevue/themes": "^4.2.5",
     "@sentry/vue": "^8.48.0",
diff --git a/src/scripts/app.ts b/src/scripts/app.ts
index dc6ec6221..82829f401 100644
--- a/src/scripts/app.ts
+++ b/src/scripts/app.ts
@@ -3,10 +3,9 @@ import {
   LGraphCanvas,
   LGraphEventMode,
   LGraphNode,
-  LiteGraph,
-  strokeShape
+  LiteGraph
 } from '@comfyorg/litegraph'
-import type { IWidget, Rect, Vector2 } from '@comfyorg/litegraph'
+import type { IWidget, Vector2 } from '@comfyorg/litegraph'
 import _ from 'lodash'
 import type { ToastMessageOptions } from 'primevue/toast'
 import { reactive } from 'vue'
@@ -572,41 +571,6 @@ export class ComfyApp {
 
       const nodeErrors = self.lastNodeErrors?.[node.id]
 
-      let color = null
-      let lineWidth = 1
-      // @ts-expect-error fixme ts strict error
-      if (node.id === +self.runningNodeId) {
-        color = '#0f0'
-      } else if (self.dragOverNode && node.id === self.dragOverNode.id) {
-        color = 'dodgerblue'
-      } else if (nodeErrors?.errors) {
-        color = 'red'
-        lineWidth = 2
-      } else if (
-        self.lastExecutionError &&
-        // @ts-expect-error fixme ts strict error
-        +self.lastExecutionError.node_id === node.id
-      ) {
-        color = '#f0f'
-        lineWidth = 2
-      }
-
-      if (color) {
-        const area: Rect = [
-          0,
-          -LiteGraph.NODE_TITLE_HEIGHT,
-          size[0],
-          size[1] + LiteGraph.NODE_TITLE_HEIGHT
-        ]
-        strokeShape(ctx, area, {
-          shape: node._shape || node.constructor.shape || LiteGraph.ROUND_SHAPE,
-          thickness: lineWidth,
-          colour: color,
-          title_height: LiteGraph.NODE_TITLE_HEIGHT,
-          collapsed: node.collapsed
-        })
-      }
-
       // @ts-expect-error fixme ts strict error
       if (self.progress && node.id === +self.runningNodeId) {
         ctx.fillStyle = 'green'
diff --git a/src/services/litegraphService.ts b/src/services/litegraphService.ts
index 91b2c68b3..fb078344d 100644
--- a/src/services/litegraphService.ts
+++ b/src/services/litegraphService.ts
@@ -46,6 +46,26 @@ export const useLitegraphService = () => {
 
       constructor(title: string) {
         super(title)
+        this.strokeStyles['running'] = function (this: LGraphNode) {
+          if (this.id == app.runningNodeId) {
+            return { color: '#0f0' }
+          }
+        }
+        this.strokeStyles['nodeError'] = function (this: LGraphNode) {
+          if (app.lastNodeErrors?.[this.id]?.errors) {
+            return { color: 'red' }
+          }
+        }
+        this.strokeStyles['dragOver'] = function (this: LGraphNode) {
+          if (app.dragOverNode?.id == this.id) {
+            return { color: 'dodgerblue' }
+          }
+        }
+        this.strokeStyles['executionError'] = function (this: LGraphNode) {
+          if (app.lastExecutionError?.node_id == this.id) {
+            return { color: '#f0f', lineWidth: 2 }
+          }
+        }
 
         const nodeMinSize = { width: 1, height: 1 }
         // Process inputs using V2 schema