From 588cfeca4bedde30a43e4c7c2a095718173ab0db Mon Sep 17 00:00:00 2001
From: Chenlei Hu <chenlei.hu@mail.utoronto.ca>
Date: Sat, 14 Sep 2024 16:01:37 +0900
Subject: [PATCH] Replace ComfyApp.runningNodeId with
 executionStore.executingNodeId (#828)

* Replace ComfyApp.runningNodeId with executionStore.executingNodeId

* nit
---
 src/scripts/app.ts               | 11 +++++++----
 src/stores/executionStore.ts     | 12 ++++++------
 tests-ui/tests/groupNode.test.ts |  3 +--
 3 files changed, 14 insertions(+), 12 deletions(-)

diff --git a/src/scripts/app.ts b/src/scripts/app.ts
index b6af221fe..11849caac 100644
--- a/src/scripts/app.ts
+++ b/src/scripts/app.ts
@@ -52,6 +52,7 @@ import { ModelStore, useModelStore } from '@/stores/modelStore'
 import type { ToastMessageOptions } from 'primevue/toast'
 import { useWorkspaceStore } from '@/stores/workspaceStateStore'
 import { ComfyLGraphNode } from '@/types/comfyLGraphNode'
+import { useExecutionStore } from '@/stores/executionStore'
 
 export const ANIM_PREVIEW_WIDGET = '$$comfy_animation_preview'
 
@@ -119,7 +120,6 @@ export class ComfyApp {
   // x, y, scale
   zoom_drag_start: [number, number, number] | null
   lastNodeErrors: any[] | null
-  runningNodeId: number | null
   lastExecutionError: { node_id: number } | null
   progress: { value: number; max: number } | null
   configuringGraph: boolean
@@ -136,6 +136,12 @@ export class ComfyApp {
   canvasContainer: HTMLElement
   menu: ComfyAppMenu
 
+  // @deprecated
+  // Use useExecutionStore().executingNodeId instead
+  get runningNodeId(): string | null {
+    return useExecutionStore().executingNodeId
+  }
+
   constructor() {
     this.vueAppReady = false
     this.ui = new ComfyUI(this)
@@ -1597,7 +1603,6 @@ export class ComfyApp {
 
     api.addEventListener('executing', ({ detail }) => {
       this.progress = null
-      this.runningNodeId = detail
       this.graph.setDirtyCanvas(true, false)
       delete this.nodePreviewImages[this.runningNodeId]
     })
@@ -1626,7 +1631,6 @@ export class ComfyApp {
     })
 
     api.addEventListener('execution_start', ({ detail }) => {
-      this.runningNodeId = null
       this.lastExecutionError = null
       // @ts-expect-error
       this.graph._nodes.forEach((node) => {
@@ -2948,7 +2952,6 @@ export class ComfyApp {
     this.nodePreviewImages = {}
     this.lastNodeErrors = null
     this.lastExecutionError = null
-    this.runningNodeId = null
   }
 
   addNodeOnGraph(
diff --git a/src/stores/executionStore.ts b/src/stores/executionStore.ts
index da1d8426a..4c158e8d7 100644
--- a/src/stores/executionStore.ts
+++ b/src/stores/executionStore.ts
@@ -10,7 +10,7 @@ export interface QueuedPrompt {
 export const useExecutionStore = defineStore('execution', () => {
   const activePromptId = ref<string | null>(null)
   const queuedPrompts = ref<Record<string, QueuedPrompt>>({})
-  const executing = ref<string | null>(null)
+  const executingNodeId = ref<string | null>(null)
 
   const activePrompt = computed(() => queuedPrompts.value[activePromptId.value])
 
@@ -65,12 +65,12 @@ export const useExecutionStore = defineStore('execution', () => {
   function handleExecuting(e: CustomEvent) {
     if (!activePrompt.value) return
 
-    if (executing.value) {
+    if (executingNodeId.value) {
       // Seems sometimes nodes that are cached fire executing but not executed
-      activePrompt.value.nodes[executing.value] = true
+      activePrompt.value.nodes[executingNodeId.value] = true
     }
-    executing.value = e.detail
-    if (!executing.value) {
+    executingNodeId.value = e.detail
+    if (!executingNodeId.value) {
       delete queuedPrompts.value[activePromptId.value]
       activePromptId.value = null
     }
@@ -104,7 +104,7 @@ export const useExecutionStore = defineStore('execution', () => {
   return {
     activePromptId,
     queuedPrompts,
-    executing,
+    executingNodeId,
     activePrompt,
     totalNodesToExecute,
     nodesExecuted,
diff --git a/tests-ui/tests/groupNode.test.ts b/tests-ui/tests/groupNode.test.ts
index a0f37388b..565d89641 100644
--- a/tests-ui/tests/groupNode.test.ts
+++ b/tests-ui/tests/groupNode.test.ts
@@ -1,3 +1,4 @@
+import { createPinia, setActivePinia } from 'pinia'
 import {
   start,
   createDefaultWorkflow,
@@ -572,7 +573,6 @@ describe('group node', () => {
       new CustomEvent('executing', { detail: `${nodes.save.id}` })
     )
     // Event should be forwarded to group node id
-    expect(+app.runningNodeId).toEqual(group.id)
     expect(group.node['imgs']).toBeFalsy()
     api.dispatchEvent(
       new CustomEvent('executed', {
@@ -613,7 +613,6 @@ describe('group node', () => {
     api.dispatchEvent(new CustomEvent('execution_start', {}))
     api.dispatchEvent(new CustomEvent('executing', { detail: `${group.id}:5` }))
     // Event should be forwarded to group node id
-    expect(+app.runningNodeId).toEqual(group.id)
     expect(group.node['imgs']).toBeFalsy()
     api.dispatchEvent(
       new CustomEvent('executed', {