Show quit confirmation dialog when sites are running

src/index.ts

@@ -12,7 +12,7 @@ import {
 import path from 'path';
 import { pathToFileURL } from 'url';
 import * as Sentry from '@sentry/electron/main';
-import { __ } from '@wordpress/i18n';
+import { __, _n, sprintf } from '@wordpress/i18n';
 import { PROTOCOL_PREFIX } from 'common/constants';
 import {
 	bumpStat,
@@ -45,8 +45,14 @@ import { renameLaunchUniquesStat } from 'src/migrations/rename-launch-uniques-st
 import { startSiteWatcher, stopSiteWatcher } from 'src/modules/cli/lib/execute-site-watch-command';
 import { updateWindowsCliVersionedPathIfNeeded } from 'src/modules/cli/lib/windows-installation-manager';
 import { setupWPServerFiles, updateWPServerFiles } from 'src/setup-wp-server-files';
-import { stopAllServersOnQuit } from 'src/site-server';
-import { loadUserData, lockAppdata, saveUserData, unlockAppdata } from 'src/storage/user-data';
+import { getRunningSiteCount, stopAllServersOnQuit } from 'src/site-server';
+import {
+	loadUserData,
+	lockAppdata,
+	saveUserData,
+	unlockAppdata,
+	updateAppdata,
+} from 'src/storage/user-data';
 import { setupUpdates } from 'src/updates';
 // eslint-disable-next-line import/order
 import packageJson from '../package.json';
@@ -95,6 +101,8 @@ const isInInstaller = require( 'electron-squirrel-startup' );
 const gotTheLock = app.requestSingleInstanceLock();
 
 let finishedInitialization = false;
+let isQuittingConfirmed = false;
+let shouldStopSitesOnQuit = true;
 
 if ( gotTheLock && ! isInInstaller ) {
 	void appBoot();
@@ -396,46 +404,99 @@ async function appBoot() {
 	} );
 
 	app.on( 'before-quit', ( event ) => {
-		if ( ! hasActiveSyncOperations() ) {
+		if ( isQuittingConfirmed ) {
 			return;
 		}
 
-		const QUIT_APP_BUTTON_INDEX = 0;
-		const CANCEL_BUTTON_INDEX = 1;
+		if ( hasActiveSyncOperations() ) {
+			const QUIT_APP_BUTTON_INDEX = 0;
+			const CANCEL_BUTTON_INDEX = 1;
+
+			const messageInformation: Pick< MessageBoxSyncOptions, 'message' | 'detail' | 'type' > =
+				hasUploadingPushOperations()
+					? {
+							message: __( 'Sync is in progress' ),
+							detail: __(
+								"There's a sync operation in progress. Quitting the app will abort that operation. Are you sure you want to quit?"
+							),
+							type: 'warning',
+					  }
+					: {
+							message: __( 'Sync will continue' ),
+							detail: __(
+								'The sync process will continue running remotely after you quit Studio. We will send you an email once it is complete.'
+							),
+							type: 'info',
+					  };
+
+			const clickedButtonIndex = dialog.showMessageBoxSync( {
+				message: messageInformation.message,
+				detail: messageInformation.detail,
+				type: messageInformation.type,
+				buttons: [ __( 'Yes, quit the app' ), __( 'No, take me back' ) ],
+				cancelId: CANCEL_BUTTON_INDEX,
+				defaultId: QUIT_APP_BUTTON_INDEX,
+			} );
 
-		const messageInformation: Pick< MessageBoxSyncOptions, 'message' | 'detail' | 'type' > =
-			hasUploadingPushOperations()
-				? {
-						message: __( 'Sync is in progress' ),
-						detail: __(
-							"There's a sync operation in progress. Quitting the app will abort that operation. Are you sure you want to quit?"
-						),
-						type: 'warning',
-				  }
-				: {
-						message: __( 'Sync will continue' ),
-						detail: __(
-							'The sync process will continue running remotely after you quit Studio. We will send you an email once it is complete.'
-						),
-						type: 'info',
-				  };
-
-		const clickedButtonIndex = dialog.showMessageBoxSync( {
-			message: messageInformation.message,
-			detail: messageInformation.detail,
-			type: messageInformation.type,
-			buttons: [ __( 'Yes, quit the app' ), __( 'No, take me back' ) ],
-			cancelId: CANCEL_BUTTON_INDEX,
-			defaultId: QUIT_APP_BUTTON_INDEX,
-		} );
+			if ( clickedButtonIndex === CANCEL_BUTTON_INDEX ) {
+				event.preventDefault();
+				return;
+			}
+		}
 
-		if ( clickedButtonIndex === CANCEL_BUTTON_INDEX ) {
+		const runningSiteCount = getRunningSiteCount();
+		if ( runningSiteCount > 0 ) {
 			event.preventDefault();
+
+			void ( async () => {
+				const userData = await loadUserData();
+
+				if ( userData.stopSitesOnQuit !== undefined ) {
+					shouldStopSitesOnQuit = userData.stopSitesOnQuit;
+					isQuittingConfirmed = true;
+					app.quit();
+					return;
+				}
+
+				const STOP_SITES_BUTTON_INDEX = 0;
+				const LEAVE_RUNNING_BUTTON_INDEX = 1;
+
+				const { response, checkboxChecked } = await dialog.showMessageBox( {
+					type: 'question',
+					message: _n( 'You have a running site', 'You have running sites', runningSiteCount ),
+					detail: sprintf(
+						_n(
+							'%d site is currently running. Do you want to stop it before quitting?',
+							'%d sites are currently running. Do you want to stop them before quitting?',
+							runningSiteCount
+						),
+						runningSiteCount
+					),
+					buttons: [ __( 'Stop sites' ), __( 'Leave running' ) ],
+					checkboxLabel: __( 'Remember my choice' ),
+					cancelId: LEAVE_RUNNING_BUTTON_INDEX,
+					defaultId: STOP_SITES_BUTTON_INDEX,
+				} );
+
+				const stopSites = response === STOP_SITES_BUTTON_INDEX;
+
+				if ( checkboxChecked ) {
+					await updateAppdata( { stopSitesOnQuit: stopSites } );
+				}
+
+				shouldStopSitesOnQuit = stopSites;
+				isQuittingConfirmed = true;
+				app.quit();
+			} )();
+
+			return;
 		}
 	} );
 
 	app.on( 'quit', () => {
-		void stopAllServersOnQuit();
+		if ( shouldStopSitesOnQuit ) {
+			void stopAllServersOnQuit();
+		}
 		stopUserDataWatcher();
 		stopSiteWatcher();
 	} );

src/site-server.ts

@@ -48,6 +48,16 @@ export async function stopAllServersOnQuit() {
 	} );
 }
 
+export function getRunningSiteCount(): number {
+	return Array.from( servers.values() ).filter( ( server ) => server.details.running ).length;
+}
+
+// Only for testing purposes
+export function __resetServersForTesting(): void {
+	servers.clear();
+	deletedServers.length = 0;
+}
+
 function getAbsoluteUrl( details: SiteDetails ): string {
 	if ( details.customDomain ) {
 		const protocol = details.enableHttps ? 'https' : 'http';

src/storage/storage-types.ts

@@ -32,6 +32,7 @@ export interface UserData {
 	preferredTerminal?: SupportedTerminal;
 	preferredEditor?: SupportedEditor;
 	betaFeatures?: BetaFeatures;
+	stopSitesOnQuit?: boolean;
 }
 
 export interface PersistedUserData extends Omit< UserData, 'sites' > {

src/storage/user-data.ts

@@ -125,6 +125,7 @@ type UserDataSafeKeys =
 	| 'onboardingCompleted'
 	| 'locale'
 	| 'promptWindowsSpeedUpResult'
+	| 'stopSitesOnQuit'
 	| 'sentryUserId'
 	| 'lastSeenVersion'
 	| 'preferredTerminal'

src/tests/site-server.test.ts

@@ -2,7 +2,7 @@
  * @jest-environment node
  */
 import { CliServerProcess } from 'src/modules/cli/lib/cli-server-process';
-import { SiteServer } from 'src/site-server';
+import { getRunningSiteCount, SiteServer, __resetServersForTesting } from 'src/site-server';
 
 // Electron's Node.js environment provides `bota`/`atob`, but Jests' does not
 jest.mock( 'common/lib/passwords' );
@@ -40,6 +40,7 @@ jest.mock( 'src/storage/user-data', () => ( {
 describe( 'SiteServer', () => {
 	beforeEach( () => {
 		jest.clearAllMocks();
+		__resetServersForTesting();
 	} );
 
 	describe( 'start', () => {
@@ -90,4 +91,59 @@ describe( 'SiteServer', () => {
 			expect( server.details.running ).toBe( true );
 		} );
 	} );
+
+	describe( 'getRunningSiteCount', () => {
+		it( 'should return 0 when no servers are registered', () => {
+			expect( getRunningSiteCount() ).toBe( 0 );
+		} );
+
+		it( 'should count only running servers', async () => {
+			const mockStart = jest.fn().mockResolvedValue( undefined );
+			const mockStop = jest.fn().mockResolvedValue( undefined );
+			( CliServerProcess as jest.Mock ).mockReturnValue( {
+				url: 'http://localhost:1234',
+				start: mockStart,
+				stop: mockStop,
+			} );
+
+			const server1 = SiteServer.register( {
+				id: 'running-site-1',
+				name: 'Running Site 1',
+				path: 'test-path-1',
+				port: 1234,
+				adminPassword: 'test-password',
+				phpVersion: '8.3',
+				running: false,
+				themeDetails: undefined,
+			} );
+			await server1.start();
+
+			SiteServer.register( {
+				id: 'stopped-site',
+				name: 'Stopped Site',
+				path: 'test-path-2',
+				port: 1235,
+				adminPassword: 'test-password',
+				phpVersion: '8.3',
+				running: false,
+				themeDetails: undefined,
+			} );
+
+			expect( getRunningSiteCount() ).toBe( 1 );
+
+			const server3 = SiteServer.register( {
+				id: 'running-site-2',
+				name: 'Running Site 2',
+				path: 'test-path-3',
+				port: 1236,
+				adminPassword: 'test-password',
+				phpVersion: '8.3',
+				running: false,
+				themeDetails: undefined,
+			} );
+			await server3.start();
+
+			expect( getRunningSiteCount() ).toBe( 2 );
+		} );
+	} );
 } );