From 3764495b63e0dd6651040ec4784014018c4b4861 Mon Sep 17 00:00:00 2001
From: huchenlei <huchenlei@proton.me>
Date: Sun, 17 Nov 2024 15:29:38 -0500
Subject: [PATCH] Move setupTray to appWindow

---
 src/main-process/appWindow.ts | 56 +++++++++++++++++++++++++++++++++-
 src/main.ts                   |  2 --
 src/tray.ts                   | 57 -----------------------------------
 3 files changed, 55 insertions(+), 60 deletions(-)
 delete mode 100644 src/tray.ts

diff --git a/src/main-process/appWindow.ts b/src/main-process/appWindow.ts
index 1e13b38a..d25cf103 100644
--- a/src/main-process/appWindow.ts
+++ b/src/main-process/appWindow.ts
@@ -1,4 +1,4 @@
-import { BrowserWindow, screen, app, shell, ipcMain } from 'electron';
+import { BrowserWindow, screen, app, shell, ipcMain, Tray, Menu } from 'electron';
 import path from 'node:path';
 import Store from 'electron-store';
 import { StoreType } from '../store';
@@ -45,6 +45,7 @@ export class AppWindow {
 
     this.setupWindowEvents();
     this.sendQueuedEventsOnReady();
+    this.setupTray();
   }
 
   public isReady(): boolean {
@@ -152,4 +153,57 @@ export class AppWindow {
       }
     });
   }
+
+  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
+    const trayImage = path.join(
+      app.isPackaged ? process.resourcesPath : './assets',
+      'UI',
+      process.platform === 'darwin' ? 'Comfy_Logo_x16_BW.png' : 'Comfy_Logo_x32.png'
+    );
+    let tray = new Tray(trayImage);
+
+    tray.setToolTip('ComfyUI');
+
+    // For Mac you can have a separate icon when you press.
+    // The current design language for Mac Eco System is White or Black icon then when you click it is in color
+    if (process.platform === 'darwin') {
+      tray.setPressedImage(path.join(app.isPackaged ? process.resourcesPath : './assets', 'UI', 'Comfy_Logo_x16.png'));
+    }
+
+    const contextMenu = Menu.buildFromTemplate([
+      {
+        label: 'Show Comfy Window',
+        click: () => {
+          this.show();
+          // Mac Only
+          if (process.platform === 'darwin') {
+            app.dock.show();
+          }
+        },
+      },
+      {
+        label: 'Quit Comfy',
+        click: () => {
+          app.quit();
+        },
+      },
+      {
+        label: 'Hide',
+        click: () => {
+          this.hide();
+          // Mac Only
+          if (process.platform === 'darwin') {
+            app.dock.hide();
+          }
+        },
+      },
+    ]);
+
+    tray.setContextMenu(contextMenu);
+
+    // If we want to make it more dynamic return tray so we can access it later
+    return tray;
+  }
 }
diff --git a/src/main.ts b/src/main.ts
index b3e2b203..a573ed10 100644
--- a/src/main.ts
+++ b/src/main.ts
@@ -2,7 +2,6 @@ import { ChildProcess } from 'node:child_process';
 import fs from 'fs';
 import axios from 'axios';
 import path from 'node:path';
-import { setupTray } from './tray';
 import { IPC_CHANNELS, SENTRY_URL_ENDPOINT, ProgressStatus } from './constants';
 import { app, dialog, ipcMain } from 'electron';
 import log from 'electron-log/main';
@@ -149,7 +148,6 @@ if (!gotTheLock) {
 
     try {
       createWindow();
-      setupTray(appWindow);
       new PathHandlers().registerHandlers();
       new AppInfoHandlers().registerHandlers();
 
diff --git a/src/tray.ts b/src/tray.ts
deleted file mode 100644
index 8f3fefbc..00000000
--- a/src/tray.ts
+++ /dev/null
@@ -1,57 +0,0 @@
-import path from 'path';
-import { Tray, Menu, app } from 'electron';
-import { AppWindow } from './main-process/appWindow';
-import { VirtualEnvironment } from './virtualEnvironment';
-
-export function setupTray(mainView: AppWindow): Tray {
-  // Set icon for the tray
-  // I think there is a way to packaged the icon in so you don't need to reference resourcesPath
-  const trayImage = path.join(
-    app.isPackaged ? process.resourcesPath : './assets',
-    'UI',
-    process.platform === 'darwin' ? 'Comfy_Logo_x16_BW.png' : 'Comfy_Logo_x32.png'
-  );
-  let tray = new Tray(trayImage);
-
-  tray.setToolTip('ComfyUI');
-
-  // For Mac you can have a separate icon when you press.
-  // The current design language for Mac Eco System is White or Black icon then when you click it is in color
-  if (process.platform === 'darwin') {
-    tray.setPressedImage(path.join(app.isPackaged ? process.resourcesPath : './assets', 'UI', 'Comfy_Logo_x16.png'));
-  }
-
-  const contextMenu = Menu.buildFromTemplate([
-    {
-      label: 'Show Comfy Window',
-      click: function () {
-        mainView.show();
-        // Mac Only
-        if (process.platform === 'darwin') {
-          app.dock.show();
-        }
-      },
-    },
-    {
-      label: 'Quit Comfy',
-      click() {
-        app.quit();
-      },
-    },
-    {
-      label: 'Hide',
-      click() {
-        mainView.hide();
-        // Mac Only
-        if (process.platform === 'darwin') {
-          app.dock.hide();
-        }
-      },
-    },
-  ]);
-
-  tray.setContextMenu(contextMenu);
-
-  // If we want to make it more dynamic return tray so we can access it later
-  return tray;
-}