From 23952d97516a7e9e00dcfd720163639f07989d73 Mon Sep 17 00:00:00 2001
From: Chenlei Hu <huchenlei@proton.me>
Date: Mon, 7 Oct 2024 19:54:00 -0400
Subject: [PATCH] Show queue front icon when shift is pressed (#1157)

* Move shiftDown state to workspaceStateStore

* Queue front state
---
 src/App.vue                                   |  9 +++++-
 src/components/actionbar/ComfyQueueButton.vue | 13 +++++++--
 src/i18n.ts                                   |  2 ++
 src/scripts/app.ts                            | 29 +++++++------------
 src/stores/workspaceStateStore.ts             |  5 +++-
 tests-ui/globalSetup.ts                       |  9 ++++++
 6 files changed, 44 insertions(+), 23 deletions(-)

diff --git a/src/App.vue b/src/App.vue
index 675f4fc87..a52f77026 100644
--- a/src/App.vue
+++ b/src/App.vue
@@ -15,8 +15,15 @@ import { useWorkspaceStore } from '@/stores/workspaceStateStore'
 import BlockUI from 'primevue/blockui'
 import ProgressSpinner from 'primevue/progressspinner'
 import GlobalDialog from '@/components/dialog/GlobalDialog.vue'
+import { useEventListener } from '@vueuse/core'
 
-const isLoading = computed<boolean>(() => useWorkspaceStore().spinner)
+const workspaceStore = useWorkspaceStore()
+const isLoading = computed<boolean>(() => workspaceStore.spinner)
+const handleKey = (e: KeyboardEvent) => {
+  workspaceStore.shiftDown = e.shiftKey
+}
+useEventListener(window, 'keydown', handleKey)
+useEventListener(window, 'keyup', handleKey)
 
 onMounted(() => {
   window['__COMFYUI_FRONTEND_VERSION__'] = config.app_version
diff --git a/src/components/actionbar/ComfyQueueButton.vue b/src/components/actionbar/ComfyQueueButton.vue
index b81667849..854a61bb5 100644
--- a/src/components/actionbar/ComfyQueueButton.vue
+++ b/src/components/actionbar/ComfyQueueButton.vue
@@ -3,13 +3,20 @@
     <SplitButton
       class="comfyui-queue-button"
       :label="activeQueueModeMenuItem.label"
-      :icon="activeQueueModeMenuItem.icon"
       severity="primary"
       @click="queuePrompt"
       :model="queueModeMenuItems"
       data-testid="queue-button"
-      v-tooltip.bottom="$t('menu.queueWorkflow')"
+      v-tooltip.bottom="
+        workspaceStore.shiftDown
+          ? $t('menu.queueWorkflowFront')
+          : $t('menu.queueWorkflow')
+      "
     >
+      <template #icon>
+        <i-lucide:list-start v-if="workspaceStore.shiftDown" />
+        <i v-else :class="activeQueueModeMenuItem.icon" />
+      </template>
       <template #item="{ item }">
         <Button
           :label="item.label"
@@ -58,7 +65,9 @@ import type { MenuItem } from 'primevue/menuitem'
 import { storeToRefs } from 'pinia'
 import { computed } from 'vue'
 import { useCommandStore } from '@/stores/commandStore'
+import { useWorkspaceStore } from '@/stores/workspaceStateStore'
 
+const workspaceStore = useWorkspaceStore()
 const queueCountStore = storeToRefs(useQueuePendingTaskCountStore())
 const { mode: queueMode } = storeToRefs(useQueueSettingsStore())
 
diff --git a/src/i18n.ts b/src/i18n.ts
index d06ce6931..b73289669 100644
--- a/src/i18n.ts
+++ b/src/i18n.ts
@@ -85,6 +85,7 @@ const messages = {
       change: 'On Change',
       changeTooltip: 'The workflow will be queued once a change is made',
       queueWorkflow: 'Queue workflow',
+      queueWorkflowFront: 'Queue workflow (Insert at Front)',
       queue: 'Queue',
       interrupt: 'Cancel current run',
       refresh: 'Refresh node definitions',
@@ -192,6 +193,7 @@ const messages = {
       change: '变动',
       changeTooltip: '工作流将会在改变后执行',
       queueWorkflow: '执行工作流',
+      queueWorkflowFront: '执行工作流 (队列首)',
       queue: '队列',
       interrupt: '取消当前任务',
       refresh: '刷新节点',
diff --git a/src/scripts/app.ts b/src/scripts/app.ts
index 17ee4c685..c9f2ee6a8 100644
--- a/src/scripts/app.ts
+++ b/src/scripts/app.ts
@@ -113,7 +113,6 @@ export class ComfyApp {
   extensionManager: ExtensionManager
   _nodeOutputs: Record<string, any>
   nodePreviewImages: Record<string, typeof Image>
-  shiftDown: boolean
   graph: LGraph
   enableWorkflowViewRestore: any
   canvas: LGraphCanvas
@@ -139,12 +138,20 @@ export class ComfyApp {
   menu: ComfyAppMenu
   bypassBgColor: string
 
-  // @deprecated
-  // Use useExecutionStore().executingNodeId instead
+  /**
+   * @deprecated Use useExecutionStore().executingNodeId instead
+   */
   get runningNodeId(): string | null {
     return useExecutionStore().executingNodeId
   }
 
+  /**
+   * @deprecated Use useWorkspaceStore().shiftDown instead
+   */
+  get shiftDown(): boolean {
+    return useWorkspaceStore().shiftDown
+  }
+
   constructor() {
     this.vueAppReady = false
     this.ui = new ComfyUI(this)
@@ -177,12 +184,6 @@ export class ComfyApp {
      * @type {Record<string, Image>}
      */
     this.nodePreviewImages = {}
-
-    /**
-     * If the shift key on the keyboard is pressed
-     * @type {boolean}
-     */
-    this.shiftDown = false
   }
 
   get nodeOutputs() {
@@ -1617,15 +1618,6 @@ export class ComfyApp {
     api.init()
   }
 
-  #addKeyboardHandler() {
-    window.addEventListener('keydown', (e) => {
-      this.shiftDown = e.shiftKey
-    })
-    window.addEventListener('keyup', (e) => {
-      this.shiftDown = e.shiftKey
-    })
-  }
-
   #addConfigureHandler() {
     const app = this
     const configure = LGraph.prototype.configure
@@ -1906,7 +1898,6 @@ export class ComfyApp {
     this.#addDropHandler()
     this.#addCopyHandler()
     this.#addPasteHandler()
-    this.#addKeyboardHandler()
     this.#addWidgetLinkHandling()
 
     await this.#invokeExtensionsAsync('setup')
diff --git a/src/stores/workspaceStateStore.ts b/src/stores/workspaceStateStore.ts
index c685bea15..483ca7132 100644
--- a/src/stores/workspaceStateStore.ts
+++ b/src/stores/workspaceStateStore.ts
@@ -7,11 +7,14 @@ import { useSidebarTabStore } from './workspace/sidebarTabStore'
 
 interface WorkspaceState {
   spinner: boolean
+  // Whether the shift key is down globally
+  shiftDown: boolean
 }
 
 export const useWorkspaceStore = defineStore('workspace', {
   state: (): WorkspaceState => ({
-    spinner: false
+    spinner: false,
+    shiftDown: false
   }),
   getters: {
     toast(): ToastManager {
diff --git a/tests-ui/globalSetup.ts b/tests-ui/globalSetup.ts
index 4217738f7..54643f774 100644
--- a/tests-ui/globalSetup.ts
+++ b/tests-ui/globalSetup.ts
@@ -45,6 +45,15 @@ module.exports = async function () {
     }
   })
 
+  jest.mock('@/stores/workspaceStateStore', () => {
+    return {
+      useWorkspaceStore: () => ({
+        shiftDown: false,
+        spinner: false
+      })
+    }
+  })
+
   jest.mock('vue-i18n', () => {
     return {
       useI18n: jest.fn()