From 3c70c1e46356ca6a18b5f9a1d7af1c8f2763d6c2 Mon Sep 17 00:00:00 2001
From: bymyself <abolkonsky.rem@gmail.com>
Date: Sun, 6 Oct 2024 09:08:16 -0700
Subject: [PATCH] Show keybinding on topbar dropdown menus (#1127)

* Show keybinding on topbar dropdown menus, resolve #1092

* Add text-muted to tailwind config

* Add Playwright test

* Preserve Primevue classes in menu item template

* Extend MenuItem

* Revert adding undo/redo to core keybindings

* Change test selector

* refactor

* Extract as component

* refactor

* nit

* fix extension API

---------

Co-authored-by: huchenlei <huchenlei@proton.me>
---
 browser_tests/menu.spec.ts                    | 10 +++
 src/components/actionbar/ComfyQueueButton.vue |  1 +
 src/components/topbar/CommandMenubar.vue      | 55 ++++++++++++
 src/components/topbar/TopMenubar.vue          | 29 +------
 src/stores/commandStore.ts                    | 55 +++++++++++-
 src/stores/menuItemStore.ts                   | 85 ++++++++-----------
 tailwind.config.js                            |  4 +
 7 files changed, 159 insertions(+), 80 deletions(-)
 create mode 100644 src/components/topbar/CommandMenubar.vue

diff --git a/browser_tests/menu.spec.ts b/browser_tests/menu.spec.ts
index 801ac3d66f..5b7561f11b 100644
--- a/browser_tests/menu.spec.ts
+++ b/browser_tests/menu.spec.ts
@@ -499,6 +499,16 @@ test.describe('Menu', () => {
       })
       expect(isTextCutoff).toBe(false)
     })
+
+    test('Displays keybinding next to item', async ({ comfyPage }) => {
+      const workflowMenuItem =
+        await comfyPage.menu.topbar.getMenuItem('Workflow')
+      await workflowMenuItem.click()
+      const exportTag = comfyPage.page.locator('.keybinding-tag', {
+        hasText: 'Ctrl + s'
+      })
+      expect(await exportTag.count()).toBe(1)
+    })
   })
 
   // Only test 'Top' to reduce test time.
diff --git a/src/components/actionbar/ComfyQueueButton.vue b/src/components/actionbar/ComfyQueueButton.vue
index c36f78deb0..b816678491 100644
--- a/src/components/actionbar/ComfyQueueButton.vue
+++ b/src/components/actionbar/ComfyQueueButton.vue
@@ -47,6 +47,7 @@
 import SplitButton from 'primevue/splitbutton'
 import Button from 'primevue/button'
 import BatchCountEdit from './BatchCountEdit.vue'
+import ButtonGroup from 'primevue/buttongroup'
 import { useI18n } from 'vue-i18n'
 import {
   AutoQueueMode,
diff --git a/src/components/topbar/CommandMenubar.vue b/src/components/topbar/CommandMenubar.vue
new file mode 100644
index 0000000000..cfa79486e7
--- /dev/null
+++ b/src/components/topbar/CommandMenubar.vue
@@ -0,0 +1,55 @@
+<template>
+  <Menubar
+    :model="items"
+    class="top-menubar border-none p-0 bg-transparent"
+    :pt="{
+      rootList: 'gap-0 flex-nowrap w-auto',
+      submenu: `dropdown-direction-${dropdownDirection}`,
+      item: 'relative'
+    }"
+  >
+    <template #item="{ item, props }">
+      <a class="p-menubar-item-link" v-bind="props.action">
+        <span v-if="item.icon" class="p-menubar-item-icon" :class="item.icon" />
+        <span class="p-menubar-item-label">{{ item.label }}</span>
+        <span
+          v-if="item?.comfyCommand?.keybinding"
+          class="ml-auto border border-surface rounded text-muted text-xs p-1 keybinding-tag"
+        >
+          {{ item.comfyCommand.keybinding.combo.toString() }}
+        </span>
+      </a>
+    </template>
+  </Menubar>
+</template>
+
+<script setup lang="ts">
+import { useMenuItemStore } from '@/stores/menuItemStore'
+import { useSettingStore } from '@/stores/settingStore'
+import Menubar from 'primevue/menubar'
+import { computed } from 'vue'
+
+const settingStore = useSettingStore()
+const dropdownDirection = computed(() =>
+  settingStore.get('Comfy.UseNewMenu') === 'Top' ? 'down' : 'up'
+)
+
+const menuItemsStore = useMenuItemStore()
+const items = menuItemsStore.menuItems
+</script>
+
+<style scoped>
+.top-menubar :deep(.p-menubar-item-link) svg {
+  display: none;
+}
+
+:deep(.p-menubar-submenu.dropdown-direction-up) {
+  @apply top-auto bottom-full flex-col-reverse;
+}
+
+.keybinding-tag {
+  background: var(--p-content-hover-background);
+  border-color: var(--p-content-border-color);
+  border-style: solid;
+}
+</style>
diff --git a/src/components/topbar/TopMenubar.vue b/src/components/topbar/TopMenubar.vue
index b6680c0252..1e79ee84c9 100644
--- a/src/components/topbar/TopMenubar.vue
+++ b/src/components/topbar/TopMenubar.vue
@@ -7,15 +7,7 @@
       :class="{ dropzone: isDropZone, 'dropzone-active': isDroppable }"
     >
       <h1 class="comfyui-logo mx-2">ComfyUI</h1>
-      <Menubar
-        :model="items"
-        class="top-menubar border-none p-0 bg-transparent"
-        :pt="{
-          rootList: 'gap-0 flex-nowrap w-auto',
-          submenu: `dropdown-direction-${dropdownDirection}`,
-          item: 'relative'
-        }"
-      />
+      <CommandMenubar />
       <Divider layout="vertical" class="mx-2" />
       <div class="flex-grow">
         <WorkflowTabs v-if="workflowTabsPosition === 'Topbar'" />
@@ -27,11 +19,10 @@
 </template>
 
 <script setup lang="ts">
-import Menubar from 'primevue/menubar'
 import Divider from 'primevue/divider'
 import WorkflowTabs from '@/components/topbar/WorkflowTabs.vue'
+import CommandMenubar from '@/components/topbar/CommandMenubar.vue'
 import Actionbar from '@/components/actionbar/ComfyActionbar.vue'
-import { useMenuItemStore } from '@/stores/menuItemStore'
 import { computed, onMounted, provide, ref } from 'vue'
 import { useSettingStore } from '@/stores/settingStore'
 import { app } from '@/scripts/app'
@@ -49,12 +40,6 @@ const teleportTarget = computed(() =>
     ? '.comfyui-body-top'
     : '.comfyui-body-bottom'
 )
-const dropdownDirection = computed(() =>
-  settingStore.get('Comfy.UseNewMenu') === 'Top' ? 'down' : 'up'
-)
-
-const menuItemsStore = useMenuItemStore()
-const items = menuItemsStore.menuItems
 
 const menuRight = ref<HTMLDivElement | null>(null)
 // Menu-right holds legacy topbar elements attached by custom scripts
@@ -105,13 +90,3 @@ eventBus.on((event: string, payload: any) => {
   cursor: default;
 }
 </style>
-
-<style>
-.top-menubar .p-menubar-item-link svg {
-  display: none;
-}
-
-.p-menubar-submenu.dropdown-direction-up {
-  @apply top-auto bottom-full flex-col-reverse;
-}
-</style>
diff --git a/src/stores/commandStore.ts b/src/stores/commandStore.ts
index 22de10ff00..7ed61e80ee 100644
--- a/src/stores/commandStore.ts
+++ b/src/stores/commandStore.ts
@@ -16,6 +16,7 @@ import { LGraphGroup } from '@comfyorg/litegraph'
 import { useTitleEditorStore } from './graphStore'
 import { useErrorHandling } from '@/hooks/errorHooks'
 import { useWorkflowStore } from './workflowStore'
+import { type KeybindingImpl, useKeybindingStore } from './keybindingStore'
 
 export interface ComfyCommand {
   id: string
@@ -24,23 +25,67 @@ export interface ComfyCommand {
   label?: string | (() => string)
   icon?: string | (() => string)
   tooltip?: string | (() => string)
+  /** Menubar item label, if different from command label */
+  menubarLabel?: string | (() => string)
   versionAdded?: string
 }
 
+export class ComfyCommandImpl implements ComfyCommand {
+  id: string
+  function: () => void | Promise<void>
+  _label?: string | (() => string)
+  _icon?: string | (() => string)
+  _tooltip?: string | (() => string)
+  _menubarLabel?: string | (() => string)
+  versionAdded?: string
+
+  constructor(command: ComfyCommand) {
+    this.id = command.id
+    this.function = command.function
+    this._label = command.label
+    this._icon = command.icon
+    this._tooltip = command.tooltip
+    this._menubarLabel = command.menubarLabel ?? command.label
+    this.versionAdded = command.versionAdded
+  }
+
+  get label() {
+    return typeof this._label === 'function' ? this._label() : this._label
+  }
+
+  get icon() {
+    return typeof this._icon === 'function' ? this._icon() : this._icon
+  }
+
+  get tooltip() {
+    return typeof this._tooltip === 'function' ? this._tooltip() : this._tooltip
+  }
+
+  get menubarLabel() {
+    return typeof this._menubarLabel === 'function'
+      ? this._menubarLabel()
+      : this._menubarLabel
+  }
+
+  get keybinding(): KeybindingImpl | null {
+    return useKeybindingStore().getKeybindingByCommandId(this.id)
+  }
+}
+
 const getTracker = () =>
   app.workflowManager.activeWorkflow?.changeTracker ?? globalTracker
 
 export const useCommandStore = defineStore('command', () => {
   const settingStore = useSettingStore()
 
-  const commandsById = ref<Record<string, ComfyCommand>>({})
+  const commandsById = ref<Record<string, ComfyCommandImpl>>({})
   const commands = computed(() => Object.values(commandsById.value))
 
   const registerCommand = (command: ComfyCommand) => {
     if (commandsById.value[command.id]) {
       console.warn(`Command ${command.id} already registered`)
     }
-    commandsById.value[command.id] = command
+    commandsById.value[command.id] = new ComfyCommandImpl(command)
   }
 
   const commandDefinitions: ComfyCommand[] = [
@@ -48,6 +93,7 @@ export const useCommandStore = defineStore('command', () => {
       id: 'Comfy.NewBlankWorkflow',
       icon: 'pi pi-plus',
       label: 'New Blank Workflow',
+      menubarLabel: 'New',
       function: () => {
         app.workflowManager.setWorkflow(null)
         app.clean()
@@ -59,6 +105,7 @@ export const useCommandStore = defineStore('command', () => {
       id: 'Comfy.OpenWorkflow',
       icon: 'pi pi-folder-open',
       label: 'Open Workflow',
+      menubarLabel: 'Open',
       function: () => {
         app.ui.loadFile()
       }
@@ -75,6 +122,7 @@ export const useCommandStore = defineStore('command', () => {
       id: 'Comfy.SaveWorkflow',
       icon: 'pi pi-save',
       label: 'Save Workflow',
+      menubarLabel: 'Save',
       function: () => {
         app.workflowManager.activeWorkflow.save()
       }
@@ -83,6 +131,7 @@ export const useCommandStore = defineStore('command', () => {
       id: 'Comfy.SaveWorkflowAs',
       icon: 'pi pi-save',
       label: 'Save Workflow As',
+      menubarLabel: 'Save As',
       function: () => {
         app.workflowManager.activeWorkflow.save(true)
       }
@@ -91,6 +140,7 @@ export const useCommandStore = defineStore('command', () => {
       id: 'Comfy.ExportWorkflow',
       icon: 'pi pi-download',
       label: 'Export Workflow',
+      menubarLabel: 'Export',
       function: () => {
         app.menu.exportWorkflow('workflow', 'workflow')
       }
@@ -99,6 +149,7 @@ export const useCommandStore = defineStore('command', () => {
       id: 'Comfy.ExportWorkflowAPI',
       icon: 'pi pi-download',
       label: 'Export Workflow (API Format)',
+      menubarLabel: 'Export (API)',
       function: () => {
         app.menu.exportWorkflow('workflow_api', 'output')
       }
diff --git a/src/stores/menuItemStore.ts b/src/stores/menuItemStore.ts
index 4c79479bc2..21b26fd8c0 100644
--- a/src/stores/menuItemStore.ts
+++ b/src/stores/menuItemStore.ts
@@ -51,61 +51,44 @@ export const useMenuItemStore = defineStore('menuItem', () => {
       commandStore.registerCommand(command)
     }
 
-    const items = commands.map(
-      (command) =>
-        ({
-          ...command,
-          command: command.function
-        }) as MenuItem
-    )
+    const items = commands
+      // Convert command to commandImpl
+      .map((command) => commandStore.getCommand(command.id))
+      .map(
+        (command) =>
+          ({
+            command: command.function,
+            label: command.menubarLabel,
+            icon: command.icon,
+            tooltip: command.tooltip,
+            comfyCommand: command
+          }) as MenuItem
+      )
     registerMenuGroup(path, items)
   }
 
-  const workflowMenuGroup: MenuItem[] = [
-    {
-      label: 'New',
-      icon: 'pi pi-plus',
-      command: () => commandStore.execute('Comfy.NewBlankWorkflow')
-    },
-    {
-      separator: true
-    },
-    {
-      label: 'Open',
-      icon: 'pi pi-folder-open',
-      command: () => commandStore.execute('Comfy.OpenWorkflow')
-    },
-    {
-      label: 'Browse Templates',
-      icon: 'pi pi-th-large',
-      command: () => commandStore.execute('Comfy.BrowseTemplates')
-    },
-    {
-      separator: true
-    },
-    {
-      label: 'Save',
-      icon: 'pi pi-save',
-      command: () => commandStore.execute('Comfy.SaveWorkflow')
-    },
-    {
-      label: 'Save As',
-      icon: 'pi pi-save',
-      command: () => commandStore.execute('Comfy.SaveWorkflowAs')
-    },
-    {
-      label: 'Export',
-      icon: 'pi pi-download',
-      command: () => commandStore.execute('Comfy.ExportWorkflow')
-    },
-    {
-      label: 'Export (API Format)',
-      icon: 'pi pi-download',
-      command: () => commandStore.execute('Comfy.ExportWorkflowAPI')
-    }
-  ]
+  registerCommands(
+    ['Workflow'],
+    [commandStore.getCommand('Comfy.NewBlankWorkflow')]
+  )
+
+  registerCommands(
+    ['Workflow'],
+    [
+      commandStore.getCommand('Comfy.OpenWorkflow'),
+      commandStore.getCommand('Comfy.BrowseTemplates')
+    ]
+  )
+  registerCommands(
+    ['Workflow'],
+    [
+      commandStore.getCommand('Comfy.SaveWorkflow'),
+      commandStore.getCommand('Comfy.SaveWorkflowAs'),
+      commandStore.getCommand('Comfy.ExportWorkflow'),
+      commandStore.getCommand('Comfy.ExportWorkflowAPI')
+    ]
+  )
 
-  registerMenuGroup(['Workflow'], workflowMenuGroup)
   registerCommands(
     ['Edit'],
     [
diff --git a/tailwind.config.js b/tailwind.config.js
index 4eaa6ac823..25d4bc7c4d 100644
--- a/tailwind.config.js
+++ b/tailwind.config.js
@@ -174,6 +174,10 @@ export default {
           900: '#7b341e',
           950: '#431407'
         }
+      },
+
+      textColor: {
+        muted: 'var(--p-text-muted-color)'
       }
     }
   },