Add executionStore (#827)

* Extract execution store

* Fix executing nodes highlight

* nit

* nit

* nit

src/App.vue

@@ -33,6 +33,7 @@ import { useQueuePendingTaskCountStore } from './stores/queueStore'
 import type { ToastMessageOptions } from 'primevue/toast'
 import { useToast } from 'primevue/usetoast'
 import { i18n } from './i18n'
+import { useExecutionStore } from './stores/executionStore'
 
 const isLoading = computed<boolean>(() => useWorkspaceStore().spinner)
 const theme = computed<string>(() =>
@@ -123,10 +124,17 @@ const onReconnected = () => {
   })
 }
 
+const executionStore = useExecutionStore()
+app.workflowManager.executionStore = executionStore
+watchEffect(() => {
+  app.menu.workflows.buttonProgress.style.width = `${executionStore.executionProgress}%`
+})
+
 onMounted(() => {
   api.addEventListener('status', onStatus)
   api.addEventListener('reconnecting', onReconnecting)
   api.addEventListener('reconnected', onReconnected)
+  executionStore.bindExecutionEvents()
   try {
     init()
   } catch (e) {
@@ -138,6 +146,7 @@ onUnmounted(() => {
   api.removeEventListener('status', onStatus)
   api.removeEventListener('reconnecting', onReconnecting)
   api.removeEventListener('reconnected', onReconnected)
+  executionStore.unbindExecutionEvents()
 })
 </script>
 

src/scripts/app.ts

@@ -1591,23 +1591,11 @@ export class ComfyApp {
     )
 
     api.addEventListener('progress', ({ detail }) => {
-      if (
-        this.workflowManager.activePrompt?.workflow &&
-        this.workflowManager.activePrompt.workflow !==
-          this.workflowManager.activeWorkflow
-      )
-        return
       this.progress = detail
       this.graph.setDirtyCanvas(true, false)
     })
 
     api.addEventListener('executing', ({ detail }) => {
-      if (
-        this.workflowManager.activePrompt?.workflow &&
-        this.workflowManager.activePrompt.workflow !==
-          this.workflowManager.activeWorkflow
-      )
-        return
       this.progress = null
       this.runningNodeId = detail
       this.graph.setDirtyCanvas(true, false)
@@ -1615,12 +1603,6 @@ export class ComfyApp {
     })
 
     api.addEventListener('executed', ({ detail }) => {
-      if (
-        this.workflowManager.activePrompt?.workflow &&
-        this.workflowManager.activePrompt.workflow !==
-          this.workflowManager.activeWorkflow
-      )
-        return
       const output = this.nodeOutputs[detail.display_node || detail.node]
       if (detail.merge && output) {
         for (const k in detail.output ?? {}) {

src/scripts/changeTracker.ts

@@ -204,7 +204,8 @@ export class ChangeTracker {
 
     // Store node outputs
     api.addEventListener('executed', ({ detail }) => {
-      const prompt = app.workflowManager.queuedPrompts[detail.prompt_id]
+      const prompt =
+        app.workflowManager.executionStore.queuedPrompts[detail.prompt_id]
       if (!prompt?.workflow) return
       const nodeOutputs = (prompt.workflow.changeTracker.nodeOutputs ??= {})
       const output = nodeOutputs[detail.node]

src/scripts/ui/menu/workflows.ts

@@ -71,17 +71,6 @@ export class ComfyWorkflowsMenu {
     })
   }
 
-  #updateProgress = () => {
-    const prompt = this.app.workflowManager.activePrompt
-    let percent = 0
-    if (this.app.workflowManager.activeWorkflow === prompt?.workflow) {
-      const total = Object.values(prompt.nodes)
-      const done = total.filter(Boolean)
-      percent = (done.length / total.length) * 100
-    }
-    this.buttonProgress.style.width = percent + '%'
-  }
-
   #updateActive = () => {
     const active = this.app.workflowManager.activeWorkflow
     this.button.tooltip = active.path
@@ -93,8 +82,6 @@ export class ComfyWorkflowsMenu {
       this.#first = false
       this.content.load()
     }
-
-    this.#updateProgress()
   }
 
   #bindEvents() {
@@ -109,10 +96,6 @@ export class ComfyWorkflowsMenu {
       this.unsaved = this.app.workflowManager.activeWorkflow.unsaved
     })
 
-    this.app.workflowManager.addEventListener('execute', (e) => {
-      this.#updateProgress()
-    })
-
     api.addEventListener('graphChanged', () => {
       this.unsaved = true
     })
@@ -371,9 +354,6 @@ export class ComfyWorkflowsContent {
       app.workflowManager.addEventListener(e, () => this.updateOpen())
     }
     this.app.workflowManager.addEventListener('rename', () => this.load())
-    this.app.workflowManager.addEventListener('execute', (e) =>
-      this.#updateActive()
-    )
   }
 
   async load() {

src/scripts/workflows.ts

@@ -17,17 +17,14 @@ export function trimJsonExt(path: string) {
 }
 
 export class ComfyWorkflowManager extends EventTarget {
-  #activePromptId: string | null = null
+  executionStore: any = null
+
   #unsavedCount = 0
   #activeWorkflow: ComfyWorkflow
 
   workflowLookup: Record<string, ComfyWorkflow> = {}
   workflows: Array<ComfyWorkflow> = []
   openWorkflows: Array<ComfyWorkflow> = []
-  queuedPrompts: Record<
-    string,
-    { workflow?: ComfyWorkflow; nodes?: Record<string, boolean> }
-  > = {}
   app: ComfyApp
 
   get activeWorkflow() {
@@ -35,62 +32,17 @@ export class ComfyWorkflowManager extends EventTarget {
   }
 
   get activePromptId() {
-    return this.#activePromptId
+    return this.executionStore?.activePromptId
   }
 
   get activePrompt() {
-    return this.queuedPrompts[this.#activePromptId]
+    return this.executionStore?.activePrompt
   }
 
   constructor(app: ComfyApp) {
     super()
     this.app = app
     ChangeTracker.init(app)
-
-    this.#bindExecutionEvents()
-  }
-
-  #bindExecutionEvents() {
-    // TODO: on reload, set active prompt based on the latest ws message
-
-    const emit = () =>
-      this.dispatchEvent(
-        new CustomEvent('execute', { detail: this.activePrompt })
-      )
-    let executing = null
-    api.addEventListener('execution_start', (e) => {
-      this.#activePromptId = e.detail.prompt_id
-
-      // This event can fire before the event is stored, so put a placeholder
-      this.queuedPrompts[this.#activePromptId] ??= { nodes: {} }
-      emit()
-    })
-    api.addEventListener('execution_cached', (e) => {
-      if (!this.activePrompt) return
-      for (const n of e.detail.nodes) {
-        this.activePrompt.nodes[n] = true
-      }
-      emit()
-    })
-    api.addEventListener('executed', (e) => {
-      if (!this.activePrompt) return
-      this.activePrompt.nodes[e.detail.node] = true
-      emit()
-    })
-    api.addEventListener('executing', (e) => {
-      if (!this.activePrompt) return
-
-      if (executing) {
-        // Seems sometimes nodes that are cached fire executing but not executed
-        this.activePrompt.nodes[executing] = true
-      }
-      executing = e.detail
-      if (!executing) {
-        delete this.queuedPrompts[this.#activePromptId]
-        this.#activePromptId = null
-      }
-      emit()
-    })
   }
 
   async loadWorkflows() {
@@ -174,15 +126,11 @@ export class ComfyWorkflowManager extends EventTarget {
   }
 
   storePrompt({ nodes, id }) {
-    this.queuedPrompts[id] ??= {}
-    this.queuedPrompts[id].nodes = {
-      ...nodes.reduce((p, n) => {
-        p[n] = false
-        return p
-      }, {}),
-      ...this.queuedPrompts[id].nodes
-    }
-    this.queuedPrompts[id].workflow = this.activeWorkflow
+    this.executionStore?.storePrompt({
+      nodes,
+      id,
+      workflow: this.activeWorkflow
+    })
   }
 
   /**

src/stores/executionStore.ts

@@ -0,0 +1,116 @@
+import { ref, computed } from 'vue'
+import { defineStore } from 'pinia'
+import { api } from '../scripts/api'
+
+export interface QueuedPrompt {
+  nodes: Record<string, boolean>
+  workflow?: any // TODO: Replace 'any' with the actual type of workflow
+}
+
+export const useExecutionStore = defineStore('execution', () => {
+  const activePromptId = ref<string | null>(null)
+  const queuedPrompts = ref<Record<string, QueuedPrompt>>({})
+  const executing = ref<string | null>(null)
+
+  const activePrompt = computed(() => queuedPrompts.value[activePromptId.value])
+
+  const totalNodesToExecute = computed(() => {
+    if (!activePrompt.value) return 0
+    return Object.values(activePrompt.value.nodes).length
+  })
+
+  const nodesExecuted = computed(() => {
+    if (!activePrompt.value) return 0
+    return Object.values(activePrompt.value.nodes).filter(Boolean).length
+  })
+
+  const executionProgress = computed(() => {
+    if (!activePrompt.value) return 0
+    const total = totalNodesToExecute.value
+    const done = nodesExecuted.value
+    return Math.round((done / total) * 100)
+  })
+
+  function bindExecutionEvents() {
+    api.addEventListener('execution_start', handleExecutionStart)
+    api.addEventListener('execution_cached', handleExecutionCached)
+    api.addEventListener('executed', handleExecuted)
+    api.addEventListener('executing', handleExecuting)
+  }
+
+  function unbindExecutionEvents() {
+    api.removeEventListener('execution_start', handleExecutionStart)
+    api.removeEventListener('execution_cached', handleExecutionCached)
+    api.removeEventListener('executed', handleExecuted)
+    api.removeEventListener('executing', handleExecuting)
+  }
+
+  function handleExecutionStart(e: CustomEvent) {
+    activePromptId.value = e.detail.prompt_id
+    queuedPrompts.value[activePromptId.value] ??= { nodes: {} }
+  }
+
+  function handleExecutionCached(e: CustomEvent) {
+    if (!activePrompt.value) return
+    for (const n of e.detail.nodes) {
+      activePrompt.value.nodes[n] = true
+    }
+  }
+
+  function handleExecuted(e: CustomEvent) {
+    if (!activePrompt.value) return
+    activePrompt.value.nodes[e.detail.node] = true
+  }
+
+  function handleExecuting(e: CustomEvent) {
+    if (!activePrompt.value) return
+
+    if (executing.value) {
+      // Seems sometimes nodes that are cached fire executing but not executed
+      activePrompt.value.nodes[executing.value] = true
+    }
+    executing.value = e.detail
+    if (!executing.value) {
+      delete queuedPrompts.value[activePromptId.value]
+      activePromptId.value = null
+    }
+  }
+
+  function storePrompt({
+    nodes,
+    id,
+    workflow
+  }: {
+    nodes: string[]
+    id: string
+    workflow: any
+  }) {
+    queuedPrompts.value[id] ??= { nodes: {} }
+    const queuedPrompt = queuedPrompts.value[id]
+    queuedPrompt.nodes = {
+      ...nodes.reduce((p, n) => {
+        p[n] = false
+        return p
+      }, {}),
+      ...queuedPrompt.nodes
+    }
+    queuedPrompt.workflow = workflow
+
+    console.debug(
+      `queued task ${id} with ${Object.values(queuedPrompt.nodes).length} nodes`
+    )
+  }
+
+  return {
+    activePromptId,
+    queuedPrompts,
+    executing,
+    activePrompt,
+    totalNodesToExecute,
+    nodesExecuted,
+    executionProgress,
+    bindExecutionEvents,
+    unbindExecutionEvents,
+    storePrompt
+  }
+})