fix(tray): avoid passing tray ref and fix update on exit during restarting (#960)

* fix(tray): avoid passing tray ref and fix update on exit

* fix: remove console.log

* test: remove tray ref

Author: samuv

main/src/auto-update.ts

@@ -11,7 +11,7 @@ import { safeTrayDestroy } from './system-tray'
 import { pollWindowReady } from './util'
 import { delay } from '../../utils/delay'
 import log from './logger'
-import { setQuittingState, setTearingDownState, getTray } from './app-state'
+import { setQuittingState, setTearingDownState } from './app-state'
 
 let pendingUpdateVersion: string | null = null
 let updateState:
@@ -79,7 +79,7 @@ async function performUpdateInstallation(releaseName: string | null) {
       log.error('[update] Error stopping ToolHive: ', error)
     }
 
-    safeTrayDestroy(getTray())
+    safeTrayDestroy()
 
     log.info('[update] all cleaned up, calling autoUpdater.quitAndInstall()...')
     autoUpdater.quitAndInstall()
@@ -89,7 +89,7 @@ async function performUpdateInstallation(releaseName: string | null) {
 
     // Attempt recovery
     try {
-      safeTrayDestroy(getTray())
+      safeTrayDestroy()
       log.info('[update] attempting app relaunch after update failure')
       // this creates a new app instance
       app.relaunch()

main/src/main.ts

@@ -104,7 +104,6 @@ import {
   setQuittingState,
   getTearingDownState,
   setTearingDownState,
-  setTray,
   getTray,
 } from './app-state'
 import type { UIMessage } from 'ai'
@@ -185,7 +184,7 @@ export async function blockQuit(source: string, event?: Electron.Event) {
     // Stop the embedded ToolHive server
     stopToolhive()
 
-    safeTrayDestroy(getTray())
+    safeTrayDestroy()
     app.quit()
   }
 }
@@ -221,10 +220,10 @@ app.whenReady().then(async () => {
 
   // Initialize tray first
   try {
-    setTray(initTray({ toolHiveIsRunning: false })) // Start with false, will update after ToolHive starts
+    initTray({ toolHiveIsRunning: false }) // Start with false, will update after ToolHive starts
     log.info('System tray initialized successfully')
     // Setup application menu
-    createApplicationMenu(getTray())
+    createApplicationMenu()
   } catch (error) {
     log.error('Failed to initialize system tray: ', error)
   }
@@ -236,7 +235,7 @@ app.whenReady().then(async () => {
   )
 
   // Start ToolHive with tray reference
-  await startToolhive(getTray() || undefined)
+  await startToolhive()
 
   // Create main window
   try {
@@ -297,7 +296,7 @@ app.whenReady().then(async () => {
     if (getTray() && process.platform !== 'win32') {
       try {
         getTray()?.destroy()
-        setTray(initTray({ toolHiveIsRunning: isToolhiveRunning() }))
+        initTray({ toolHiveIsRunning: isToolhiveRunning() })
       } catch (error) {
         log.error('Failed to update tray after theme change: ', error)
       }
@@ -357,7 +356,7 @@ app.on('will-quit', (e) => blockQuit('will-quit', e))
       }
     } finally {
       stopToolhive()
-      safeTrayDestroy(getTray())
+      safeTrayDestroy()
       process.exit(0)
     }
   })
@@ -391,11 +390,9 @@ ipcMain.handle('get-auto-launch-status', () => getAutoLaunchStatus())
 ipcMain.handle('set-auto-launch', (_event, enabled: boolean) => {
   setAutoLaunch(enabled)
   // Update tray menu if exists
-  if (getTray()) {
-    updateTrayStatus(getTray()!, isToolhiveRunning())
-  }
+  updateTrayStatus(isToolhiveRunning())
   // Update menu
-  createApplicationMenu(getTray())
+  createApplicationMenu()
   return getAutoLaunchStatus()
 })
 
@@ -464,7 +461,7 @@ ipcMain.handle('check-container-engine', async () => {
 
 ipcMain.handle('restart-toolhive', async () => {
   try {
-    await restartToolhive(getTray() || undefined)
+    await restartToolhive()
     return { success: true }
   } catch (error) {
     log.error('Failed to restart ToolHive: ', error)

main/src/menu.ts

@@ -3,10 +3,7 @@ import { getAutoLaunchStatus, setAutoLaunch } from './auto-launch'
 import { updateTrayStatus } from './system-tray'
 import log from './logger'
 
-function createAutoLaunchItem(
-  accelerator: string,
-  trayRef: Electron.Tray | null
-) {
+function createAutoLaunchItem(accelerator: string) {
   return {
     label: 'Start on Login',
     type: 'checkbox' as const,
@@ -16,18 +13,16 @@ function createAutoLaunchItem(
       try {
         const currentStatus = getAutoLaunchStatus()
         setAutoLaunch(!currentStatus)
-        if (trayRef) {
-          updateTrayStatus(trayRef, true)
-        }
-        createApplicationMenu(trayRef)
+        updateTrayStatus(true)
+        createApplicationMenu()
       } catch (error) {
         log.error('Failed to toggle auto-launch: ', error)
       }
     },
   }
 }
 
-export function createApplicationMenu(trayRef: Electron.Tray | null) {
+export function createApplicationMenu() {
   const isMac = process.platform === 'darwin'
   const defaultMenu = Menu.getApplicationMenu()?.items ?? []
 
@@ -57,7 +52,7 @@ export function createApplicationMenu(trayRef: Electron.Tray | null) {
       submenu: [
         { role: 'about' as const },
         { type: 'separator' as const },
-        createAutoLaunchItem(isMac ? 'Cmd+L' : 'Ctrl+L', trayRef),
+        createAutoLaunchItem(isMac ? 'Cmd+L' : 'Ctrl+L'),
         { type: 'separator' as const },
         ...(isMac
           ? [

main/src/system-tray.ts

@@ -6,9 +6,11 @@ import log from './logger'
 import { getAppVersion } from './util'
 import { hideWindow, showWindow, showInDock } from './dock-utils'
 import { showMainWindow, sendToMainWindowRenderer } from './main-window'
+import { getTray, setTray } from './app-state'
 
 // Safe tray destruction with error handling
-export function safeTrayDestroy(tray: Tray | null) {
+export function safeTrayDestroy() {
+  const tray = getTray()
   try {
     if (tray && !tray.isDestroyed()) {
       tray.destroy()
@@ -62,13 +64,10 @@ function getIcon(): Electron.NativeImage {
 ///////////////////////////////////////////////////
 
 const createTrayWithSetup =
-  (setupFn: (tray: Tray, toolHiveIsRunning: boolean) => void) =>
+  (setupFn: (toolHiveIsRunning: boolean) => void) =>
   (toolHiveIsRunning: boolean) => {
     try {
-      const tray = new Tray(getIcon())
-
-      setupFn(tray, toolHiveIsRunning)
-      return tray
+      return setupFn(toolHiveIsRunning)
     } catch (error) {
       log.error('Failed to create tray: ', error)
       throw error
@@ -111,10 +110,7 @@ const showWindowWithFocus = (window: BrowserWindow) => {
   bringToFrontOnWindows(window)
 }
 
-const handleStartOnLogin = async (
-  currentTray: Tray,
-  toolHiveIsRunning: boolean
-) => {
+const handleStartOnLogin = async (toolHiveIsRunning: boolean) => {
   const currentStatus = getAutoLaunchStatus()
 
   try {
@@ -127,10 +123,10 @@ const handleStartOnLogin = async (
     }
 
     // Update the tray menu to reflect the new state
-    setupTrayMenu(currentTray, toolHiveIsRunning)
+    setupTrayMenu(toolHiveIsRunning)
 
     // Update the application menu to reflect the new state
-    createApplicationMenu(currentTray)
+    createApplicationMenu()
   } catch (error) {
     log.error('Failed to toggle auto-launch: ', error)
   }
@@ -153,14 +149,14 @@ const getCurrentAppVersion = () => {
   }
 }
 
-const startOnLoginMenu = (currentTray: Tray, toolHiveIsRunning: boolean) => {
+const startOnLoginMenu = (toolHiveIsRunning: boolean) => {
   const isStartOnLogin = getAutoLaunchStatus()
   return {
     label: 'Start on login',
     checked: isStartOnLogin,
     accelerator: 'CmdOrCtrl+L',
     type: 'checkbox' as const,
-    click: () => handleStartOnLogin(currentTray, toolHiveIsRunning),
+    click: () => handleStartOnLogin(toolHiveIsRunning),
   }
 }
 
@@ -195,11 +191,11 @@ const createQuitMenuItem = () => ({
 
 const createSeparator = () => ({ type: 'separator' as const })
 
-const createMenuTemplate = (currentTray: Tray, toolHiveIsRunning: boolean) => [
+const createMenuTemplate = (toolHiveIsRunning: boolean) => [
   createStatusMenuItem(toolHiveIsRunning),
   getCurrentAppVersion(),
   createSeparator(),
-  startOnLoginMenu(currentTray, toolHiveIsRunning),
+  startOnLoginMenu(toolHiveIsRunning),
   createSeparator(),
   createShowMenuItem(),
   createHideMenuItem(),
@@ -233,8 +229,15 @@ const createClickHandler = () => {
   }
 }
 
-function setupTrayMenu(tray: Tray, toolHiveIsRunning: boolean) {
-  const menuTemplate = createMenuTemplate(tray, toolHiveIsRunning)
+// if tray is not there let's create a new one
+function setupTrayMenu(toolHiveIsRunning: boolean) {
+  let tray = getTray()
+  if (!tray || tray.isDestroyed()) {
+    tray = new Tray(getIcon())
+    setTray(tray)
+  }
+
+  const menuTemplate = createMenuTemplate(toolHiveIsRunning)
   const contextMenu = Menu.buildFromTemplate(menuTemplate)
 
   tray.setToolTip('ToolHive')
@@ -247,11 +250,12 @@ function setupTrayMenu(tray: Tray, toolHiveIsRunning: boolean) {
 }
 
 // Function to update tray status without recreating the entire tray
-export const updateTrayStatus = (tray: Tray, toolHiveIsRunning: boolean) => {
+export const updateTrayStatus = (toolHiveIsRunning: boolean) => {
+  const tray = getTray()
   if (!tray || tray.isDestroyed()) return
 
   try {
-    setupTrayMenu(tray, toolHiveIsRunning)
+    setupTrayMenu(toolHiveIsRunning)
   } catch (error) {
     log.error('Failed to update tray status: ', error)
   }

main/src/tests/auto-update.test.ts

@@ -237,7 +237,7 @@ describe('auto-update', () => {
         expect(vi.mocked(dialog.showMessageBox)).toHaveBeenCalled()
         expect(vi.mocked(stopAllServers)).toHaveBeenCalled()
         expect(vi.mocked(stopToolhive)).toHaveBeenCalled()
-        expect(vi.mocked(safeTrayDestroy)).toHaveBeenCalledWith(mockTray)
+        expect(vi.mocked(safeTrayDestroy)).toHaveBeenCalled()
       })
 
       it('handles update-downloaded event with user clicking later', async () => {
@@ -681,7 +681,7 @@ describe('auto-update', () => {
       // Should call graceful exit sequence
       expect(vi.mocked(stopAllServers)).toHaveBeenCalled()
       expect(vi.mocked(stopToolhive)).toHaveBeenCalled()
-      expect(vi.mocked(safeTrayDestroy)).toHaveBeenCalledWith(mockTray)
+      expect(vi.mocked(safeTrayDestroy)).toHaveBeenCalled()
     })
 
     it('handles server shutdown failure during update', async () => {
@@ -789,9 +789,7 @@ describe('auto-update', () => {
 
       await updatePromise
 
-      expect(vi.mocked(safeTrayDestroy)).toHaveBeenCalledWith(
-        mockTrayWithDestroy
-      )
+      expect(vi.mocked(safeTrayDestroy)).toHaveBeenCalled()
     })
 
     it('respects event listener removal during update', async () => {