From 82df3cf8115b9a70dcc47072cb658838bb2a2e90 Mon Sep 17 00:00:00 2001
From: filtered <176114999+webfiltered@users.noreply.github.com>
Date: Thu, 12 Dec 2024 05:36:32 +1100
Subject: [PATCH] Add native context menu (#463)

* Add interop to open system context menu

* Allow default context menu pos (at cursor)

* Add text context menu

- Replaces earlier incarnation of system context menu
---
 src/constants.ts                    |  1 +
 src/main-process/appWindow.ts       | 25 +++++++++++++++++++++++--
 src/main-process/comfyDesktopApp.ts |  7 +++++--
 src/preload.ts                      | 14 ++++++++++++++
 4 files changed, 43 insertions(+), 4 deletions(-)

diff --git a/src/constants.ts b/src/constants.ts
index 74bc5d99..7bcfac71 100644
--- a/src/constants.ts
+++ b/src/constants.ts
@@ -30,6 +30,7 @@ export const IPC_CHANNELS = {
   VALIDATE_COMFYUI_SOURCE: 'validate-comfyui-source',
   SHOW_DIRECTORY_PICKER: 'show-directory-picker',
   INSTALL_COMFYUI: 'install-comfyui',
+  SHOW_CONTEXT_MENU: 'show-context-menu',
 } as const;
 
 export enum ProgressStatus {
diff --git a/src/main-process/appWindow.ts b/src/main-process/appWindow.ts
index a7b36a09..5dcb9ff6 100644
--- a/src/main-process/appWindow.ts
+++ b/src/main-process/appWindow.ts
@@ -1,4 +1,4 @@
-import { BrowserWindow, screen, app, shell, ipcMain, Tray, Menu, dialog, MenuItem } from 'electron';
+import { BrowserWindow, screen, app, shell, ipcMain, Tray, Menu, dialog, MenuItem, type Point } from 'electron';
 import path from 'node:path';
 import Store from 'electron-store';
 import { AppWindowSettings } from '../store';
@@ -6,6 +6,7 @@ import log from 'electron-log/main';
 import { IPC_CHANNELS, ProgressStatus, ServerArgs } from '../constants';
 import { getAppResourcesPath } from '../install/resourcePaths';
 import { DesktopConfig } from '../store/desktopConfig';
+import type { ElectronContextMenuOptions } from '../preload';
 
 /**
  * Creates a single application window that displays the renderer and encapsulates all the logic for sending messages to the renderer.
@@ -13,9 +14,14 @@ import { DesktopConfig } from '../store/desktopConfig';
  */
 export class AppWindow {
   private window: BrowserWindow;
+  /** Volatile store containing window config - saves window state between launches. */
   private store: Store<AppWindowSettings>;
   private messageQueue: Array<{ channel: string; data: any }> = [];
   private rendererReady: boolean = false;
+  /** The application menu. */
+  private menu: Electron.Menu | null;
+  /** The "edit" menu - cut/copy/paste etc. */
+  private editMenu?: Menu;
 
   public constructor() {
     const installed = DesktopConfig.store.get('installState') === 'installed';
@@ -55,7 +61,8 @@ export class AppWindow {
     this.setupAppEvents();
     this.sendQueuedEventsOnReady();
     this.setupTray();
-    this.buildMenu();
+    this.menu = this.buildMenu();
+    this.buildTextMenu();
   }
 
   public isReady(): boolean {
@@ -224,6 +231,14 @@ export class AppWindow {
     });
   }
 
+  showSystemContextMenu(options?: ElectronContextMenuOptions): void {
+    if (options?.type === 'text') {
+      this.editMenu?.popup(options.pos);
+    } else {
+      this.menu?.popup(options?.pos);
+    }
+  }
+
   setupTray() {
     // Set icon for the tray
     // I think there is a way to packaged the icon in so you don't need to reference resourcesPath
@@ -277,6 +292,11 @@ export class AppWindow {
     return tray;
   }
 
+  buildTextMenu() {
+    // Electron bug - strongly typed to the incorrect case.
+    this.editMenu = Menu.getApplicationMenu()?.items.find((x) => x.role?.toLowerCase() === 'editmenu')?.submenu;
+  }
+
   buildMenu() {
     const menu = Menu.getApplicationMenu();
     if (menu) {
@@ -306,5 +326,6 @@ export class AppWindow {
         Menu.setApplicationMenu(menu);
       }
     }
+    return menu;
   }
 }
diff --git a/src/main-process/comfyDesktopApp.ts b/src/main-process/comfyDesktopApp.ts
index c32f8f48..cdc8097f 100644
--- a/src/main-process/comfyDesktopApp.ts
+++ b/src/main-process/comfyDesktopApp.ts
@@ -1,4 +1,4 @@
-import { app, dialog, ipcMain } from 'electron';
+import { app, dialog, ipcMain, type Point } from 'electron';
 import log from 'electron-log/main';
 import * as Sentry from '@sentry/electron/main';
 import { graphics } from 'systeminformation';
@@ -9,7 +9,7 @@ import { AppWindow } from './appWindow';
 import { ComfyServer } from './comfyServer';
 import { ComfyServerConfig } from '../config/comfyServerConfig';
 import fs from 'fs';
-import { InstallOptions } from '../preload';
+import { InstallOptions, type ElectronContextMenuOptions } from '../preload';
 import path from 'path';
 import { getModelsDirectory, validateHardware } from '../utils';
 import { DownloadManager } from '../models/DownloadManager';
@@ -92,6 +92,9 @@ export class ComfyDesktopApp {
   }
 
   registerIPCHandlers(): void {
+    ipcMain.on(IPC_CHANNELS.SHOW_CONTEXT_MENU, (_event, options?: ElectronContextMenuOptions) => {
+      this.appWindow.showSystemContextMenu(options);
+    });
     ipcMain.on(IPC_CHANNELS.OPEN_DEV_TOOLS, () => {
       this.appWindow.openDevTools();
     });
diff --git a/src/preload.ts b/src/preload.ts
index f44ac90e..0b181631 100644
--- a/src/preload.ts
+++ b/src/preload.ts
@@ -35,6 +35,11 @@ export interface DownloadProgressUpdate {
   message?: string;
 }
 
+export interface ElectronContextMenuOptions {
+  type: 'system' | 'text' | 'image';
+  pos?: Electron.Point;
+}
+
 const electronAPI = {
   /**
    * Callback for progress updates from the main process for starting ComfyUI.
@@ -215,6 +220,15 @@ const electronAPI = {
   installComfyUI: (installOptions: InstallOptions) => {
     ipcRenderer.send(IPC_CHANNELS.INSTALL_COMFYUI, installOptions);
   },
+  /**
+   * Opens native context menus.
+   *
+   * {@link ElectronContextMenuOptions} contains the various options to control the menu type.
+   * @param options Define which type of menu to use, position, etc.
+   */
+  showContextMenu: (options?: ElectronContextMenuOptions): void => {
+    return ipcRenderer.send(IPC_CHANNELS.SHOW_CONTEXT_MENU, options);
+  },
 } as const;
 
 export type ElectronAPI = typeof electronAPI;