fix(quit): wait for ComfyUI to exit before quitting or installing an update

src/main/index.ts

@@ -52,9 +52,12 @@ import { waitForPort, COMFY_BOOT_TIMEOUT_MS } from './lib/process'
 import {
   clearQuitReason,
   isQuitInProgress,
+  relaunchIfQuitting,
+  scheduleRelaunch,
   setQuitReason,
   setSessionEnding
 } from './lib/quit-state'
+import { holdQuit, isQuitHeld } from './lib/quitWait'
 import { showUpdateInstallSplash } from './lib/updateSplash'
 import type { InstallationRecord } from './installations'
 import {
@@ -793,7 +796,7 @@ ipcMain.handle('app:relaunch', () => {
     tray.destroy()
     tray = null
   }
-  app.relaunch()
+  scheduleRelaunch(() => app.relaunch())
   app.quit()
 })
 
@@ -1342,6 +1345,8 @@ if (app.isPackaged && !app.requestSingleInstanceLock()) {
 
   if (app.isPackaged) {
     app.on('second-instance', () => {
+      // Quitting: the window would close with the app, so relaunch after the quit instead.
+      if (relaunchIfQuitting(() => app.relaunch())) return
       // OS-level "open another instance" attempt - focus an existing
       // host window (chooser or install-backed) instead of stacking
       // a duplicate. Queued until startup recovery settles.
@@ -1480,7 +1485,6 @@ if (app.isPackaged && !app.requestSingleInstanceLock()) {
       sessionStartProps: { prior_session_unclean: priorSessionUnclean }
     })
     mainTelemetry.setConsentState(initialConsent)
-    mainTelemetry.installAppHooks()
 
     // installation_id is an event/person property, never a PostHog identity.
     const existingInstallation = hasCompletedFirstLaunch() || hasPersistedDeviceId()
@@ -2417,6 +2421,12 @@ if (app.isPackaged && !app.requestSingleInstanceLock()) {
   })
 
   app.on('before-quit', (event) => {
+    // A quit already waiting for ComfyUI to exit (closing the windows below quits again, via
+    // window-all-closed): hold this one too, or the process exits mid-wait.
+    if (isQuitHeld()) {
+      event.preventDefault()
+      return
+    }
     // Template models are still downloading in the background: quitting pauses
     // them (they restore as resumable rows in Downloads next launch), but the
     // template setup itself stops until the user resumes them. Confirm once and
@@ -2437,33 +2447,40 @@ if (app.isPackaged && !app.requestSingleInstanceLock()) {
         return
       }
     }
-    if (!isQuitInProgress()) {
-      setQuitReason('user-quit')
-      ipc.cancelAll()
-      for (const [, entry] of comfyWindows) {
-        if (!entry.window.isDestroyed()) entry.window.destroy()
+    try {
+      if (!isQuitInProgress()) {
+        setQuitReason('user-quit')
+        ipc.cancelAll()
+        for (const [, entry] of comfyWindows) {
+          if (!entry.window.isDestroyed()) entry.window.destroy()
+        }
+        comfyWindows.clear()
+        // Pop-out terminal/logs windows live outside `comfyWindows`; close them
+        // here too and kill their shared shells so no window or PTY child lingers.
+        closeAllPopouts()
+        disposeAllTerminals()
+        if (tray) {
+          tray.destroy()
+          tray = null
+        }
       }
-      comfyWindows.clear()
-      // Pop-out terminal/logs windows live outside `comfyWindows`; close them
-      // here too and kill their shared shells so no window or PTY child lingers.
-      closeAllPopouts()
-      disposeAllTerminals()
-      if (tray) {
-        tray.destroy()
-        tray = null
+      updater.recordProcessExit()
+      if (_stopPeriodicReleaseChecks) {
+        _stopPeriodicReleaseChecks()
+        _stopPeriodicReleaseChecks = null
       }
+      // Persist any pending last-active-surface write so the next boot can
+      // restore it. Synchronous: the app exits without awaiting promises, so an
+      // async write would be torn down mid-flight and lose a just-made change.
+      flushLastSessionSync()
+      flushOperationOutput()
+      cleanupTempDownloads()
+    } finally {
+      // Hold the quit until the ComfyUI processes being stopped have exited (bounded) and
+      // telemetry has drained, then quit again; that re-issued quit passes through. In a
+      // finally: a throw above must not cost the wait or the drain.
+      holdQuit(event, { drain: mainTelemetry.drainForQuit, quit: () => app.quit() })
     }
-    updater.recordProcessExit()
-    if (_stopPeriodicReleaseChecks) {
-      _stopPeriodicReleaseChecks()
-      _stopPeriodicReleaseChecks = null
-    }
-    // Persist any pending last-active-surface write so the next boot can
-    // restore it. Synchronous: the app exits without awaiting promises, so an
-    // async write would be torn down mid-flight and lose a just-made change.
-    flushLastSessionSync()
-    flushOperationOutput()
-    cleanupTempDownloads()
   })
 
   // Deferred-quit suspension for managed model downloads: stop each active

src/main/lib/ipc/launchCancel.test.ts

@@ -26,6 +26,7 @@ import {
   _test_addRunningSession,
   _test_clearRunningSessions
 } from './shared'
+import { _resetQuitWaitForTest, waitForExitWork } from '../quitWait'
 
 const INSTALL = 'install-under-test'
 
@@ -46,6 +47,36 @@ describe('cancelAll', () => {
 
     expect(flushTelemetry).toHaveBeenCalledOnce()
   })
+
+  it('aborts a launch still in its pre-spawn prep, and a quit waits for its teardown', async () => {
+    _resetQuitWaitForTest()
+    const launch = _beginLaunch(INSTALL)
+    cancelAll()
+    expect(launch.abort.signal.aborted).toBe(true)
+    let waited = false
+    const waiting = waitForExitWork().then(() => {
+      waited = true
+    })
+    await new Promise((r) => setTimeout(r, 20))
+    expect(waited).toBe(false)
+    _endLaunch(INSTALL, launch)
+    await waiting
+  })
+
+  it('a launch whose slot a relaunch took still settles, so later quits do not wait for it', async () => {
+    _resetQuitWaitForTest()
+    // A restart while the first launch is parked after registration (the template gate).
+    const first = _beginLaunch(INSTALL)
+    const second = _beginLaunch(INSTALL)
+    _endLaunch(INSTALL, second)
+    _endLaunch(INSTALL, first)
+    let waited = false
+    void waitForExitWork().then(() => {
+      waited = true
+    })
+    await new Promise((r) => setTimeout(r, 20))
+    expect(waited).toBe(true)
+  })
 })
 
 describe('launch tracking (_beginLaunch/_endLaunch/_hasActiveLaunch)', () => {

src/main/lib/ipc/sessionActions/launch.test.ts

@@ -106,6 +106,8 @@ const launchHarness = vi.hoisted(() => ({
   portLockPid: null as null | number,
   /** What the mocked `killProcessTree` reports: false = the tree outlived the kill wait. */
   killExits: true,
+  /** What `checkRebootMarker` answers (a Manager restart request); null = the real check. */
+  rebootMarker: null as null | boolean,
   campaigns: { registry: [] as unknown[], answers: new Map<string, unknown>() },
   campaignRecords: {} as Record<string, Record<string, { epoch: number; enrolledAt: number }>>,
   recordWrites: [] as Array<[string, string, number]>,
@@ -244,7 +246,9 @@ vi.mock('../shared', async (importOriginal) => {
         : actual.getProcessInfo(...args),
     // Never let a test reach the real one: the fake child's pid is invented, and killing it
     // would signal whatever real process happens to hold that pid.
-    killProcessTree: async () => ({ exited: launchHarness.killExits, waitMs: 0 })
+    killProcessTree: async () => ({ exited: launchHarness.killExits, waitMs: 0 }),
+    checkRebootMarker: (...args: Parameters<typeof actual.checkRebootMarker>) =>
+      launchHarness.rebootMarker ?? actual.checkRebootMarker(...args)
   }
 })
 
@@ -376,6 +380,8 @@ import {
   _resolveLaunchMode,
   _resolvePortConflictPolicy
 } from './launch'
+import { _resetQuitWaitForTest, waitForExitWork } from '../../quitWait'
+import { clearQuitReason, setQuitReason } from '../../quit-state'
 import * as assetsTapModule from '../../assetsTap'
 import * as agentTapModule from '../../agentTap'
 import {
@@ -588,6 +594,61 @@ describe('onProcessTerminated', () => {
   })
 })
 
+describe('onProcessTerminated holds a quit for the exit bookkeeping', () => {
+  beforeEach(() => {
+    _resetQuitWaitForTest()
+  })
+
+  it('from the exit until the callback has finished', async () => {
+    const proc = new EventEmitter() as unknown as ChildProcess
+    let finish!: () => void
+    onProcessTerminated(proc, () => new Promise<void>((resolve) => (finish = resolve)))
+    proc.emit('exit', 0, null)
+    proc.emit('close', 0, null)
+    let waited = false
+    const waiting = waitForExitWork().then(() => {
+      waited = true
+    })
+    await new Promise((r) => setTimeout(r, 20))
+    expect(waited).toBe(false)
+    finish()
+    await waiting
+  })
+
+  it('and lets it go when the callback throws', async () => {
+    const proc = new EventEmitter() as unknown as ChildProcess
+    const consoleError = vi.spyOn(console, 'error').mockImplementation(() => {})
+    try {
+      onProcessTerminated(proc, () => {
+        throw new Error('sync failure')
+      })
+      proc.emit('exit', 1, null)
+      proc.emit('close', 1, null)
+      await waitForExitWork()
+    } finally {
+      consoleError.mockRestore()
+    }
+  })
+})
+
+describe('handleLaunch during a quit', () => {
+  afterEach(() => {
+    clearQuitReason()
+  })
+
+  it('refuses to start a ComfyUI while Desktop is quitting or installing an update', async () => {
+    setQuitReason('update-install')
+    const res = await handleLaunch({
+      event: { sender: { send: vi.fn() } } as unknown as Electron.IpcMainInvokeEvent,
+      installationId: 'launch-during-quit',
+      inst: installOf('not-a-real-source'),
+      actionData: {}
+    })
+    expect(res).toEqual({ ok: false, cancelled: true })
+    expect(_operationAborts.has('launch-during-quit')).toBe(false)
+  })
+})
+
 describe('_cleanupFailedLaunchSetup', () => {
   const INSTALL = 'cleanup-under-test'
   const PORT = 59_311
@@ -1424,6 +1485,7 @@ describe('core beta report placement', () => {
     launchHarness.commitShas = []
     launchHarness.duringResourceAcquire = null
     launchHarness.waitForPort = null
+    launchHarness.rebootMarker = null
     // Both halves of the activation-notice state: the in-process pending queue and the
     // persisted announced list, which the real settings module keeps in this run's temp
     // app dir. Without the reset, the first test to launch spends the notice for the rest.
@@ -1476,6 +1538,32 @@ describe('core beta report placement', () => {
     return proc
   }
 
+  it('does not respawn on a Manager restart request while Desktop is quitting', async () => {
+    let spawns = 0
+    let child: FakeChild | null = null
+    launchHarness.spawn = () => {
+      spawns++
+      child = fakeChild()
+      return child
+    }
+    // The port-wait path: the one that wires the long-lived exit handler.
+    launchHarness.launchCommand = { ...launchHarness.launchCommand, skipPortWait: false }
+    launchHarness.waitForPort = async () => {}
+    const res = await handleLaunch(ctxFor('harness-quit-no-respawn'))
+    expect(res.ok).toBe(true)
+    _resetQuitWaitForTest()
+    launchHarness.rebootMarker = true
+    setQuitReason('user-quit')
+    try {
+      child!.emit('exit', 0, null)
+      child!.emit('close', 0, null)
+      await waitForExitWork()
+      expect(spawns).toBe(1)
+    } finally {
+      clearQuitReason()
+    }
+  })
+
   describe('campaigns', () => {
     const KEY = 'desktop_core_beta_agent'
     const agentAnswer = (variant: string): unknown =>
@@ -2961,7 +3049,7 @@ describe('Performance Test guardrail', () => {
   let started: InstanceStartedCallbackInfo[] = []
   let installDir = ''
   let bootLogs: unknown[] = []
-  const launches: [string, { abort: AbortController }][] = []
+  const launches: [string, ReturnType<typeof _beginLaunch>][] = []
   // Sessions earlier tests in this file left registered would all count as running here.
   let leftRunning: [string, SessionInfo][] = []
   let leftStopping: string[] = []

src/main/lib/ipc/sessionActions/launch.ts

@@ -117,6 +117,8 @@ import {
 import type { AgentStartOutcome, OverridePins } from '../../agentRequirementsOverride'
 import { recoverInterruptedComfyOp } from '../../opMarker'
 import { waitLaunchSpawnHold } from '../../e2eOverrides'
+import { trackExitWork } from '../../quitWait'
+import { isQuitInProgress } from '../../quit-state'
 import {
   holderIsInstall,
   resolvePriorProcess,
@@ -623,22 +625,32 @@ export function onProcessTerminated(
 ): void {
   let finished = false
   let fallbackTimer: NodeJS.Timeout | undefined
+  // From `'exit'` until the callback has run, a quit waits for it: the record clear and
+  // `comfyui.exited` happen there.
+  let handled!: () => void
+  const handledPromise = new Promise<void>((resolve) => {
+    handled = resolve
+  })
 
   const finish = (code: number | null, signal: NodeJS.Signals | null, pipesHeld: boolean): void => {
     if (finished) return
     finished = true
     if (fallbackTimer) clearTimeout(fallbackTimer)
     try {
-      void Promise.resolve(callback(code, signal, { pipesHeld })).catch((err) => {
-        console.error('Process termination callback failed:', err)
-      })
+      void Promise.resolve(callback(code, signal, { pipesHeld }))
+        .catch((err) => {
+          console.error('Process termination callback failed:', err)
+        })
+        .finally(handled)
     } catch (err) {
       console.error('Process termination callback failed:', err)
+      handled()
     }
   }
 
   proc.once('close', (code, signal) => finish(code, signal, false))
   proc.once('exit', (code, signal) => {
+    trackExitWork(handledPromise)
     if (finished) return
     fallbackTimer = setTimeout(() => finish(code, signal, true), PROCESS_CLOSE_GRACE_MS)
     fallbackTimer.unref()
@@ -983,6 +995,9 @@ export async function handleLaunch(ctx: ActionContext): Promise<ActionResult> {
   if (_runningSessions.has(sessionId)) {
     return { ok: false, message: i18n.t('errors.alreadyRunning') }
   }
+  // Quitting or installing an update: a ComfyUI started now would outlive the wait for the
+  // ones being stopped.
+  if (isQuitInProgress()) return { ok: false, cancelled: true }
   // No `_hasActiveLaunch` here: this guard, `_beginLaunch`, and runLaunch's
   // `_operationAborts.set` all run in one synchronous stretch, so a second
   // launch can never slip between them. Checking it would instead reject a
@@ -2713,7 +2728,10 @@ async function runLaunch(
         rebootModelCheckAbort = null
       }
 
-      if (pendingModelFolderRelaunch || checkRebootMarker(sessionPath)) {
+      // Never respawn while quitting: the new ComfyUI would outlive the quit's wait. The marker
+      // is still consumed, as before.
+      const relaunch = pendingModelFolderRelaunch || checkRebootMarker(sessionPath)
+      if (relaunch && !isQuitInProgress()) {
         const isModelRelaunch = pendingModelFolderRelaunch
         pendingModelFolderRelaunch = false
         if (!isModelRelaunch) {