From 57a4cb9036d2f762b6f6c5c1ebb30e1465216735 Mon Sep 17 00:00:00 2001
From: Chenlei Hu <huchenlei@proton.me>
Date: Fri, 4 Oct 2024 16:28:08 -0400
Subject: [PATCH] Add default toast error handling for command execution
 (#1106)

* Error handling execute command

* Cleanup

* Add playwright test

* Mock i18n in jest test

* Reduce test func timeout
---
 browser_tests/ComfyPage.ts                    | 30 ++++++++++++
 browser_tests/commands.spec.ts                | 49 +++++++++++++++++++
 src/components/actionbar/ComfyActionbar.vue   | 13 ++---
 src/components/graph/GraphCanvasMenu.vue      | 13 ++---
 .../sidebar/tabs/QueueSidebarTab.vue          |  4 +-
 .../sidebar/tabs/WorkflowsSidebarTab.vue      | 10 ++--
 src/extensions/core/keybinds.ts               |  2 +-
 src/hooks/errorHooks.ts                       |  2 +-
 src/stores/commandStore.ts                    | 16 +++---
 src/stores/menuItemStore.ts                   | 14 +++---
 src/stores/workspaceStateStore.ts             |  6 +++
 src/types/extensionTypes.ts                   |  5 ++
 tests-ui/globalSetup.ts                       |  6 +++
 13 files changed, 126 insertions(+), 44 deletions(-)
 create mode 100644 browser_tests/commands.spec.ts

diff --git a/browser_tests/ComfyPage.ts b/browser_tests/ComfyPage.ts
index eedd5885df..5b08c456e5 100644
--- a/browser_tests/ComfyPage.ts
+++ b/browser_tests/ComfyPage.ts
@@ -485,6 +485,36 @@ export class ComfyPage {
     return `./browser_tests/assets/${fileName}`
   }
 
+  async executeCommand(commandId: string) {
+    await this.page.evaluate((id: string) => {
+      return window['app'].extensionManager.command.execute(id)
+    }, commandId)
+  }
+
+  async registerCommand(
+    commandId: string,
+    command: (() => void) | (() => Promise<void>)
+  ) {
+    await this.page.evaluate(
+      ({ commandId, commandStr }) => {
+        const app = window['app']
+        const randomSuffix = Math.random().toString(36).substring(2, 8)
+        const extensionName = `TestExtension_${randomSuffix}`
+
+        app.registerExtension({
+          name: extensionName,
+          commands: [
+            {
+              id: commandId,
+              function: eval(commandStr)
+            }
+          ]
+        })
+      },
+      { commandId, commandStr: command.toString() }
+    )
+  }
+
   async registerKeybinding(keyCombo: KeyCombo, command: () => void) {
     await this.page.evaluate(
       ({ keyCombo, commandStr }) => {
diff --git a/browser_tests/commands.spec.ts b/browser_tests/commands.spec.ts
new file mode 100644
index 0000000000..7ee3744c99
--- /dev/null
+++ b/browser_tests/commands.spec.ts
@@ -0,0 +1,49 @@
+import { expect } from '@playwright/test'
+import { comfyPageFixture as test } from './ComfyPage'
+
+test.describe('Keybindings', () => {
+  test('Should execute command', async ({ comfyPage }) => {
+    await comfyPage.registerCommand('TestCommand', () => {
+      window['foo'] = true
+    })
+
+    await comfyPage.executeCommand('TestCommand')
+    expect(await comfyPage.page.evaluate(() => window['foo'])).toBe(true)
+  })
+
+  test('Should execute async command', async ({ comfyPage }) => {
+    await comfyPage.registerCommand('TestCommand', async () => {
+      await new Promise<void>((resolve) =>
+        setTimeout(() => {
+          window['foo'] = true
+          resolve()
+        }, 5)
+      )
+    })
+
+    await comfyPage.executeCommand('TestCommand')
+    expect(await comfyPage.page.evaluate(() => window['foo'])).toBe(true)
+  })
+
+  test('Should handle command errors', async ({ comfyPage }) => {
+    await comfyPage.registerCommand('TestCommand', () => {
+      throw new Error('Test error')
+    })
+
+    await comfyPage.executeCommand('TestCommand')
+    await expect(comfyPage.page.locator('.p-toast')).toBeVisible()
+  })
+
+  test('Should handle async command errors', async ({ comfyPage }) => {
+    await comfyPage.registerCommand('TestCommand', async () => {
+      await new Promise<void>((resolve, reject) =>
+        setTimeout(() => {
+          reject(new Error('Test error'))
+        }, 5)
+      )
+    })
+
+    await comfyPage.executeCommand('TestCommand')
+    await expect(comfyPage.page.locator('.p-toast')).toBeVisible()
+  })
+})
diff --git a/src/components/actionbar/ComfyActionbar.vue b/src/components/actionbar/ComfyActionbar.vue
index 6f4ad95783..fb3f8fab8f 100644
--- a/src/components/actionbar/ComfyActionbar.vue
+++ b/src/components/actionbar/ComfyActionbar.vue
@@ -37,7 +37,7 @@
             :severity="executingPrompt ? 'danger' : 'secondary'"
             :disabled="!executingPrompt"
             text
-            @click="() => commandStore.getCommandFunction('Comfy.Interrupt')()"
+            @click="() => commandStore.execute('Comfy.Interrupt')"
           >
           </Button>
           <Button
@@ -46,9 +46,7 @@
             :severity="hasPendingTasks ? 'danger' : 'secondary'"
             :disabled="!hasPendingTasks"
             text
-            @click="
-              () => commandStore.getCommandFunction('Comfy.ClearPendingTasks')()
-            "
+            @click="() => commandStore.execute('Comfy.ClearPendingTasks')"
           />
         </ButtonGroup>
       </div>
@@ -59,10 +57,7 @@
           icon="pi pi-refresh"
           severity="secondary"
           text
-          @click="
-            () =>
-              commandStore.getCommandFunction('Comfy.RefreshNodeDefinitions')()
-          "
+          @click="() => commandStore.execute('Comfy.RefreshNodeDefinitions')"
         />
       </ButtonGroup>
     </div>
@@ -143,7 +138,7 @@ const hasPendingTasks = computed(() => queueCountStore.count.value > 1)
 
 const queuePrompt = (e: MouseEvent) => {
   const commandId = e.shiftKey ? 'Comfy.QueuePromptFront' : 'Comfy.QueuePrompt'
-  commandStore.getCommandFunction(commandId)()
+  commandStore.execute(commandId)
 }
 
 const panelRef = ref<HTMLElement | null>(null)
diff --git a/src/components/graph/GraphCanvasMenu.vue b/src/components/graph/GraphCanvasMenu.vue
index 73ede6972d..3f42fe8b2d 100644
--- a/src/components/graph/GraphCanvasMenu.vue
+++ b/src/components/graph/GraphCanvasMenu.vue
@@ -20,7 +20,7 @@
       severity="secondary"
       icon="pi pi-expand"
       v-tooltip.left="t('graphCanvasMenu.resetView')"
-      @click="() => commandStore.getCommandFunction('Comfy.Canvas.ResetView')()"
+      @click="() => commandStore.execute('Comfy.Canvas.ResetView')"
     />
     <Button
       severity="secondary"
@@ -29,9 +29,7 @@
           'graphCanvasMenu.' + (canvasStore.readOnly ? 'panMode' : 'selectMode')
         ) + ' (Space)'
       "
-      @click="
-        () => commandStore.getCommandFunction('Comfy.Canvas.ToggleLock')()
-      "
+      @click="() => commandStore.execute('Comfy.Canvas.ToggleLock')"
     >
       <template #icon>
         <i-material-symbols:pan-tool-outline v-if="canvasStore.readOnly" />
@@ -42,10 +40,7 @@
       severity="secondary"
       :icon="linkHidden ? 'pi pi-eye-slash' : 'pi pi-eye'"
       v-tooltip.left="t('graphCanvasMenu.toggleLinkVisibility')"
-      @click="
-        () =>
-          commandStore.getCommandFunction('Comfy.Canvas.ToggleLinkVisibility')()
-      "
+      @click="() => commandStore.execute('Comfy.Canvas.ToggleLinkVisibility')"
       data-testid="toggle-link-visibility-button"
     />
   </ButtonGroup>
@@ -73,7 +68,7 @@ const linkHidden = computed(
 let interval: number | null = null
 const repeat = (command: string) => {
   if (interval) return
-  const cmd = commandStore.getCommandFunction(command)
+  const cmd = () => commandStore.execute(command)
   cmd()
   interval = window.setInterval(cmd, 100)
 }
diff --git a/src/components/sidebar/tabs/QueueSidebarTab.vue b/src/components/sidebar/tabs/QueueSidebarTab.vue
index d3c0336cb5..765ead9f56 100644
--- a/src/components/sidebar/tabs/QueueSidebarTab.vue
+++ b/src/components/sidebar/tabs/QueueSidebarTab.vue
@@ -36,9 +36,7 @@
           icon="pi pi-stop"
           severity="danger"
           text
-          @click="
-            () => commandStore.getCommandFunction('Comfy.ClearPendingTasks')()
-          "
+          @click="() => commandStore.execute('Comfy.ClearPendingTasks')"
           v-tooltip.bottom="$t('sideToolbar.queueTab.clearPendingTasks')"
         />
         <Button
diff --git a/src/components/sidebar/tabs/WorkflowsSidebarTab.vue b/src/components/sidebar/tabs/WorkflowsSidebarTab.vue
index e847360ba7..3506be4f43 100644
--- a/src/components/sidebar/tabs/WorkflowsSidebarTab.vue
+++ b/src/components/sidebar/tabs/WorkflowsSidebarTab.vue
@@ -6,24 +6,20 @@
         icon="pi pi-th-large"
         v-tooltip="$t('sideToolbar.browseTemplates')"
         text
-        @click="
-          () => commandStore.getCommandFunction('Comfy.BrowseTemplates')()
-        "
+        @click="() => commandStore.execute('Comfy.BrowseTemplates')"
       />
       <Button
         class="open-workflow-button"
         icon="pi pi-folder-open"
         v-tooltip="$t('sideToolbar.openWorkflow')"
         text
-        @click="() => commandStore.getCommandFunction('Comfy.OpenWorkflow')()"
+        @click="() => commandStore.execute('Comfy.OpenWorkflow')"
       />
       <Button
         class="new-blank-workflow-button"
         icon="pi pi-plus"
         v-tooltip="$t('sideToolbar.newBlankWorkflow')"
-        @click="
-          () => commandStore.getCommandFunction('Comfy.NewBlankWorkflow')()
-        "
+        @click="() => commandStore.execute('Comfy.NewBlankWorkflow')"
         text
       />
     </template>
diff --git a/src/extensions/core/keybinds.ts b/src/extensions/core/keybinds.ts
index 6b8ab4bdba..7cfdc5a45a 100644
--- a/src/extensions/core/keybinds.ts
+++ b/src/extensions/core/keybinds.ts
@@ -31,7 +31,7 @@ app.registerExtension({
       const commandStore = useCommandStore()
       const keybinding = keybindingStore.getKeybinding(keyCombo)
       if (keybinding) {
-        await commandStore.getCommandFunction(keybinding.commandId)()
+        await commandStore.execute(keybinding.commandId)
         event.preventDefault()
         return
       }
diff --git a/src/hooks/errorHooks.ts b/src/hooks/errorHooks.ts
index b45492cd88..9d5762edea 100644
--- a/src/hooks/errorHooks.ts
+++ b/src/hooks/errorHooks.ts
@@ -26,7 +26,7 @@ export function useErrorHandling() {
 
   const wrapWithErrorHandlingAsync =
     (
-      action: (...args: any[]) => Promise<any>,
+      action: ((...args: any[]) => Promise<any>) | ((...args: any[]) => any),
       errorHandler?: (error: any) => void
     ) =>
     async (...args: any[]) => {
diff --git a/src/stores/commandStore.ts b/src/stores/commandStore.ts
index 918a7bebda..9e0e072e98 100644
--- a/src/stores/commandStore.ts
+++ b/src/stores/commandStore.ts
@@ -15,6 +15,7 @@ import { ComfyExtension } from '@/types/comfy'
 import { useWorkspaceStore } from './workspaceStateStore'
 import { LGraphGroup } from '@comfyorg/litegraph'
 import { useTitleEditorStore } from './graphStore'
+import { useErrorHandling } from '@/hooks/errorHooks'
 
 export interface ComfyCommand {
   id: string
@@ -312,18 +313,20 @@ export const useCommandStore = defineStore('command', () => {
   ]
 
   commandDefinitions.forEach(registerCommand)
-  const getCommandFunction = (command: string) => {
-    return commandsById.value[command]?.function ?? (() => {})
-  }
-
   const getCommand = (command: string) => {
     return commandsById.value[command]
   }
 
-  const execute = (commandId: string) => {
+  const { wrapWithErrorHandlingAsync } = useErrorHandling()
+  const execute = async (
+    commandId: string,
+    errorHandler?: (error: any) => void
+  ) => {
     const command = getCommand(commandId)
     if (command) {
-      command.function()
+      await wrapWithErrorHandlingAsync(command.function, errorHandler)()
+    } else {
+      throw new Error(`Command ${commandId} not found`)
     }
   }
 
@@ -343,7 +346,6 @@ export const useCommandStore = defineStore('command', () => {
     commands,
     execute,
     getCommand,
-    getCommandFunction,
     registerCommand,
     isRegistered,
     loadExtensionCommands
diff --git a/src/stores/menuItemStore.ts b/src/stores/menuItemStore.ts
index e339917086..4c79479bc2 100644
--- a/src/stores/menuItemStore.ts
+++ b/src/stores/menuItemStore.ts
@@ -65,7 +65,7 @@ export const useMenuItemStore = defineStore('menuItem', () => {
     {
       label: 'New',
       icon: 'pi pi-plus',
-      command: commandStore.getCommandFunction('Comfy.NewBlankWorkflow')
+      command: () => commandStore.execute('Comfy.NewBlankWorkflow')
     },
     {
       separator: true
@@ -73,12 +73,12 @@ export const useMenuItemStore = defineStore('menuItem', () => {
     {
       label: 'Open',
       icon: 'pi pi-folder-open',
-      command: commandStore.getCommandFunction('Comfy.OpenWorkflow')
+      command: () => commandStore.execute('Comfy.OpenWorkflow')
     },
     {
       label: 'Browse Templates',
       icon: 'pi pi-th-large',
-      command: commandStore.getCommandFunction('Comfy.BrowseTemplates')
+      command: () => commandStore.execute('Comfy.BrowseTemplates')
     },
     {
       separator: true
@@ -86,22 +86,22 @@ export const useMenuItemStore = defineStore('menuItem', () => {
     {
       label: 'Save',
       icon: 'pi pi-save',
-      command: commandStore.getCommandFunction('Comfy.SaveWorkflow')
+      command: () => commandStore.execute('Comfy.SaveWorkflow')
     },
     {
       label: 'Save As',
       icon: 'pi pi-save',
-      command: commandStore.getCommandFunction('Comfy.SaveWorkflowAs')
+      command: () => commandStore.execute('Comfy.SaveWorkflowAs')
     },
     {
       label: 'Export',
       icon: 'pi pi-download',
-      command: commandStore.getCommandFunction('Comfy.ExportWorkflow')
+      command: () => commandStore.execute('Comfy.ExportWorkflow')
     },
     {
       label: 'Export (API Format)',
       icon: 'pi pi-download',
-      command: commandStore.getCommandFunction('Comfy.ExportWorkflowAPI')
+      command: () => commandStore.execute('Comfy.ExportWorkflowAPI')
     }
   ]
 
diff --git a/src/stores/workspaceStateStore.ts b/src/stores/workspaceStateStore.ts
index 04f259570e..0a145e315d 100644
--- a/src/stores/workspaceStateStore.ts
+++ b/src/stores/workspaceStateStore.ts
@@ -3,6 +3,7 @@ import { defineStore } from 'pinia'
 import { useToastStore } from './toastStore'
 import { useQueueSettingsStore } from './queueStore'
 import { useMenuItemStore } from './menuItemStore'
+import { useCommandStore } from './commandStore'
 
 interface WorkspaceState {
   spinner: boolean
@@ -27,6 +28,11 @@ export const useWorkspaceStore = defineStore('workspace', {
       return {
         registerTopbarCommands: useMenuItemStore().registerCommands
       }
+    },
+    command() {
+      return {
+        execute: useCommandStore().execute
+      }
     }
   },
   actions: {
diff --git a/src/types/extensionTypes.ts b/src/types/extensionTypes.ts
index 7693198d0e..1e71f6b11f 100644
--- a/src/types/extensionTypes.ts
+++ b/src/types/extensionTypes.ts
@@ -39,4 +39,9 @@ export interface ExtensionManager {
 
   // Toast
   toast: ToastManager
+  command: CommandManager
+}
+
+export interface CommandManager {
+  execute(command: string, errorHandler?: (error: any) => void): void
 }
diff --git a/tests-ui/globalSetup.ts b/tests-ui/globalSetup.ts
index b9c81d2877..5ad73025f5 100644
--- a/tests-ui/globalSetup.ts
+++ b/tests-ui/globalSetup.ts
@@ -27,4 +27,10 @@ module.exports = async function () {
         })
     }
   })
+
+  jest.mock('vue-i18n', () => {
+    return {
+      useI18n: jest.fn()
+    }
+  })
 }