From f55842bb72362d0c42089c2e59301d4d928932ba Mon Sep 17 00:00:00 2001
From: huchenlei <chenlei.hu@mail.utoronto.ca>
Date: Sun, 15 Sep 2024 10:19:15 +0900
Subject: [PATCH] Add workflows store to manage workflows

---
 src/App.vue                 |  4 +++
 src/scripts/workflows.ts    | 63 +++++++++++++++----------------------
 src/stores/workflowStore.ts | 13 ++++++++
 src/utils/formatUtil.ts     |  4 +--
 4 files changed, 44 insertions(+), 40 deletions(-)
 create mode 100644 src/stores/workflowStore.ts

diff --git a/src/App.vue b/src/App.vue
index b0f8755fb..e89fcb5c5 100644
--- a/src/App.vue
+++ b/src/App.vue
@@ -34,6 +34,7 @@ import type { ToastMessageOptions } from 'primevue/toast'
 import { useToast } from 'primevue/usetoast'
 import { i18n } from './i18n'
 import { useExecutionStore } from './stores/executionStore'
+import { useWorkflowStore } from './stores/workflowStore'
 
 const isLoading = computed<boolean>(() => useWorkspaceStore().spinner)
 const theme = computed<string>(() =>
@@ -130,6 +131,9 @@ watchEffect(() => {
   app.menu.workflows.buttonProgress.style.width = `${executionStore.executionProgress}%`
 })
 
+const workflowStore = useWorkflowStore()
+app.workflowManager.workflowStore = workflowStore
+
 onMounted(() => {
   api.addEventListener('status', onStatus)
   api.addEventListener('reconnecting', onReconnecting)
diff --git a/src/scripts/workflows.ts b/src/scripts/workflows.ts
index bb7087220..f7a429698 100644
--- a/src/scripts/workflows.ts
+++ b/src/scripts/workflows.ts
@@ -8,12 +8,11 @@ import { appendJsonExt, trimJsonExt } from '@/utils/formatUtil'
 
 export class ComfyWorkflowManager extends EventTarget {
   executionStore: any = null
+  workflowStore: any = null
 
   #unsavedCount = 0
   #activeWorkflow: ComfyWorkflow
 
-  workflowLookup: Record<string, ComfyWorkflow> = {}
-  workflows: Array<ComfyWorkflow> = []
   openWorkflows: Array<ComfyWorkflow> = []
   app: ComfyApp
 
@@ -29,6 +28,14 @@ export class ComfyWorkflowManager extends EventTarget {
     return this.executionStore?.activePrompt
   }
 
+  get workflows(): ComfyWorkflow[] {
+    return this.workflowStore?.workflows || []
+  }
+
+  get workflowLookup(): Record<string, ComfyWorkflow> {
+    return this.workflowStore?.workflowLookup || {}
+  }
+
   constructor(app: ComfyApp) {
     super()
     this.app = app
@@ -47,8 +54,8 @@ export class ComfyWorkflowManager extends EventTarget {
         favorites = new Set()
       }
 
-      const workflows = (await api.listUserData('workflows', true, true)).map(
-        (w) => {
+      ;(await api.listUserData('workflows', true, true)).forEach(
+        (w: string[]) => {
           let workflow = this.workflowLookup[w[0]]
           if (!workflow) {
             workflow = new ComfyWorkflow(
@@ -59,14 +66,10 @@ export class ComfyWorkflowManager extends EventTarget {
             )
             this.workflowLookup[workflow.path] = workflow
           }
-          return workflow
         }
       )
-
-      this.workflows = workflows
     } catch (error) {
       alert('Error loading workflows: ' + (error.message ?? error))
-      this.workflows = []
     }
   }
 
@@ -167,30 +170,14 @@ export class ComfyWorkflowManager extends EventTarget {
 }
 
 export class ComfyWorkflow {
-  #name
-  #path
-  #pathParts
-  #isFavorite = false
+  name: string | null = null
+  path: string | null = null
+  pathParts: string[] | null = null
+  isFavorite: boolean = false
   changeTracker: ChangeTracker | null = null
-  unsaved = false
+  unsaved: boolean = false
   manager: ComfyWorkflowManager
 
-  get name() {
-    return this.#name
-  }
-
-  get path() {
-    return this.#path
-  }
-
-  get pathParts() {
-    return this.#pathParts
-  }
-
-  get isFavorite() {
-    return this.#isFavorite
-  }
-
   get isOpen() {
     return !!this.changeTracker
   }
@@ -204,15 +191,15 @@ export class ComfyWorkflow {
     this.manager = manager
     if (pathParts) {
       this.#updatePath(path, pathParts)
-      this.#isFavorite = isFavorite
+      this.isFavorite = isFavorite
     } else {
-      this.#name = path
+      this.name = path
       this.unsaved = true
     }
   }
 
   #updatePath(path: string, pathParts: string[]) {
-    this.#path = path
+    this.path = path
 
     if (!pathParts) {
       if (!path.includes('\\')) {
@@ -222,8 +209,8 @@ export class ComfyWorkflow {
       }
     }
 
-    this.#pathParts = pathParts
-    this.#name = trimJsonExt(pathParts[pathParts.length - 1])
+    this.pathParts = pathParts
+    this.name = trimJsonExt(pathParts[pathParts.length - 1])
   }
 
   async getWorkflowData() {
@@ -266,8 +253,8 @@ export class ComfyWorkflow {
 
   async favorite(value: boolean) {
     try {
-      if (this.#isFavorite === value) return
-      this.#isFavorite = value
+      if (this.isFavorite === value) return
+      this.isFavorite = value
       await this.manager.saveWorkflowMetadata()
       this.manager.dispatchEvent(new CustomEvent('favorite', { detail: this }))
     } catch (error) {
@@ -352,8 +339,8 @@ export class ComfyWorkflow {
     }
 
     this.unsaved = true
-    this.#path = null
-    this.#pathParts = null
+    this.path = null
+    this.pathParts = null
     this.manager.workflows.splice(this.manager.workflows.indexOf(this), 1)
     this.manager.dispatchEvent(new CustomEvent('delete', { detail: this }))
   }
diff --git a/src/stores/workflowStore.ts b/src/stores/workflowStore.ts
new file mode 100644
index 000000000..5c1b8395e
--- /dev/null
+++ b/src/stores/workflowStore.ts
@@ -0,0 +1,13 @@
+import { ref, computed } from 'vue'
+import { defineStore } from 'pinia'
+import { ComfyWorkflow } from '@/scripts/workflows'
+
+export const useWorkflowStore = defineStore('workflow', () => {
+  const workflowLookup = ref<Record<string, ComfyWorkflow>>({})
+  const workflows = computed(() => Object.values(workflowLookup.value))
+
+  return {
+    workflows,
+    workflowLookup
+  }
+})
diff --git a/src/utils/formatUtil.ts b/src/utils/formatUtil.ts
index 6bec35cec..a593ab731 100644
--- a/src/utils/formatUtil.ts
+++ b/src/utils/formatUtil.ts
@@ -30,6 +30,6 @@ export function appendJsonExt(path: string) {
   return path
 }
 
-export function trimJsonExt(path: string) {
-  return path.replace(/\.json$/, '')
+export function trimJsonExt(path?: string) {
+  return path?.replace(/\.json$/, '')
 }