Generate site thumbnail when server starts (Automattic#128)

* Prototype screenshot generation

* Each screenshot has a fresh browser session

* Screenshot captured and cached each time server starts

* Hide adminbar in thumbnails

* Wait a short amount of time before taking screenshot

* Taking screenshots doesn't prevent UI from reporting server is ready

* Thumbnail cache deleted when site is deleted from Studio

---------

Co-authored-by: Wojtek Naruniec <wojtek@naruniec.me>

src/constants.ts

@@ -1,4 +1,6 @@
 export const MAIN_MIN_WIDTH = 900;
 export const MAIN_MIN_HEIGHT = 600;
+export const SCREENSHOT_WIDTH = 1040;
+export const SCREENSHOT_HEIGHT = 1248;
 export const LIMIT_OF_ZIP_SITES_PER_USER = 10;
 export const AUTO_UPDATE_INTERVAL_MS = 60 * 60 * 1000;

src/screenshot-window.ts

@@ -0,0 +1,28 @@
+import crypto from 'crypto';
+import { BrowserWindow, session } from 'electron';
+import { SCREENSHOT_HEIGHT, SCREENSHOT_WIDTH } from './constants';
+
+export function createScreenshotWindow( captureUrl: string ) {
+	const newSession = session.fromPartition( crypto.randomUUID() );
+
+	const window = new BrowserWindow( {
+		height: SCREENSHOT_HEIGHT,
+		width: SCREENSHOT_WIDTH,
+		show: false,
+		webPreferences: { session: newSession },
+	} );
+
+	const finishedLoading = new Promise< void >( ( resolve ) => {
+		window.webContents.on( 'did-finish-load', () => resolve() );
+	} );
+
+	window.loadURL( captureUrl );
+
+	const waitForCapture = async () => {
+		await finishedLoading;
+		await new Promise( ( resolve ) => setTimeout( resolve, 500 ) );
+		return window.webContents.capturePage();
+	};
+
+	return { window, waitForCapture };
+}

src/site-server.ts

@@ -1,8 +1,12 @@
 import { app } from 'electron';
+import fs from 'fs/promises';
 import nodePath from 'path';
+import * as Sentry from '@sentry/electron/main';
 import { getWpNowConfig, startServer, type WPNowServer } from '../vendor/wp-now/src';
 import { pathExists, recursiveCopyDirectory, isEmptyDir } from './lib/fs-utils';
 import { portFinder } from './lib/port-finder';
+import { createScreenshotWindow } from './screenshot-window';
+import { getSiteThumbnailPath } from './storage/paths';
 
 const servers = new Map< string, SiteServer >();
 
@@ -55,6 +59,8 @@ export class SiteServer {
 	}
 
 	async delete() {
+		const thumbnailPath = getSiteThumbnailPath( this.details.id );
+		await fs.unlink( thumbnailPath );
 		await this.stop();
 		servers.delete( this.details.id );
 	}
@@ -85,6 +91,8 @@ export class SiteServer {
 			port: this.server.options.port,
 			running: true,
 		};
+
+		await this.updateCachedThumbnail();
 	}
 
 	updateSiteDetails( site: SiteDetails ) {
@@ -107,4 +115,24 @@ export class SiteServer {
 		const { running, url, ...rest } = this.details;
 		this.details = { running: false, ...rest };
 	}
+
+	async updateCachedThumbnail() {
+		if ( ! this.details.running ) {
+			throw new Error( 'Cannot update thumbnail for a stopped server' );
+		}
+
+		const captureUrl = new URL( '/?studio-hide-adminbar', this.details.url );
+		const { window, waitForCapture } = createScreenshotWindow( captureUrl.href );
+
+		const outPath = getSiteThumbnailPath( this.details.id );
+		const outDir = nodePath.dirname( outPath );
+
+		// Continue taking the screenshot asynchronously so we don't prevent the
+		// UI from showing the server is now available.
+		fs.mkdir( outDir, { recursive: true } )
+			.then( waitForCapture )
+			.then( ( image ) => fs.writeFile( outPath, image.toPNG() ) )
+			.catch( Sentry.captureException )
+			.finally( () => window.destroy() );
+	}
 }

src/storage/paths.ts

@@ -12,3 +12,8 @@ export function getServerFilesPath(): string {
 }
 
 export const DEFAULT_SITE_PATH = path.join( app?.getPath( 'home' ) || '', 'Studio' );
+
+export function getSiteThumbnailPath( siteId: string ): string {
+	const appDataPath = app.getPath( 'appData' );
+	return path.join( appDataPath, app.getName(), 'thumbnails', `${ siteId }.png` );
+}

vendor/wp-now/src/download.ts

@@ -232,6 +232,16 @@ export async function downloadMuPlugins(customMuPluginsPath = '') {
 		}`
 	);
 
+	fs.writeFile(
+		path.join(muPluginsPath, '0-thumbnails.php'),
+		`<?php
+		// Facilitates the taking of screenshots to be used as thumbnails.
+		if ( isset( $_GET['studio-hide-adminbar'] ) ) {
+			add_filter( 'show_admin_bar', '__return_false' );
+		}
+		`
+	)
+
 	fs.writeFile(
 		path.join(muPluginsPath, '0-sqlite.php'),
 		`<?php