From 8c100a6fe6cc652c3799d1622e12c2c969f30510 Mon Sep 17 00:00:00 2001
From: Bjorn Lu <bjornlu.dev@gmail.com>
Date: Sat, 14 Jan 2023 16:51:01 +0800
Subject: [PATCH] Handle server restart from Vite plugins (#5849)

---
 .changeset/poor-ladybugs-thank.md             |  5 ++
 packages/astro/src/core/dev/restart.ts        | 64 +++++++++++--------
 packages/astro/test/units/dev/restart.test.js | 31 +++++++++
 3 files changed, 72 insertions(+), 28 deletions(-)
 create mode 100644 .changeset/poor-ladybugs-thank.md

diff --git a/.changeset/poor-ladybugs-thank.md b/.changeset/poor-ladybugs-thank.md
new file mode 100644
index 000000000..109ce43a9
--- /dev/null
+++ b/.changeset/poor-ladybugs-thank.md
@@ -0,0 +1,5 @@
+---
+'astro': patch
+---
+
+Handle server restart from Vite plugins
diff --git a/packages/astro/src/core/dev/restart.ts b/packages/astro/src/core/dev/restart.ts
index 9ded41325..9c71b7aa3 100644
--- a/packages/astro/src/core/dev/restart.ts
+++ b/packages/astro/src/core/dev/restart.ts
@@ -142,35 +142,39 @@ export async function createContainerWithAutomaticRestart({
 		},
 	};
 
-	function handleServerRestart(logMsg: string) {
+	async function handleServerRestart(logMsg: string) {
 		// eslint-disable-next-line @typescript-eslint/no-shadow
 		const container = restart.container;
-		return async function (changedFile: string) {
-			if (shouldRestartContainer(container, changedFile)) {
-				const { container: newContainer, error } = await restartContainer({
-					beforeRestart,
-					container,
-					flags,
-					logMsg,
-					async handleConfigError(err) {
-						// Send an error message to the client if one is connected.
-						await handleConfigError(err);
-						container.viteServer.ws.send({
-							type: 'error',
-							err: {
-								message: err.message,
-								stack: err.stack || '',
-							},
-						});
+		const { container: newContainer, error } = await restartContainer({
+			beforeRestart,
+			container,
+			flags,
+			logMsg,
+			async handleConfigError(err) {
+				// Send an error message to the client if one is connected.
+				await handleConfigError(err);
+				container.viteServer.ws.send({
+					type: 'error',
+					err: {
+						message: err.message,
+						stack: err.stack || '',
 					},
 				});
-				restart.container = newContainer;
-				// Add new watches because this is a new container with a new Vite server
-				addWatches();
-				resolveRestart(error);
-				restartComplete = new Promise<Error | null>((resolve) => {
-					resolveRestart = resolve;
-				});
+			},
+		});
+		restart.container = newContainer;
+		// Add new watches because this is a new container with a new Vite server
+		addWatches();
+		resolveRestart(error);
+		restartComplete = new Promise<Error | null>((resolve) => {
+			resolveRestart = resolve;
+		});
+	}
+
+	function handleChangeRestart(logMsg: string) {
+		return async function (changedFile: string) {
+			if (shouldRestartContainer(restart.container, changedFile)) {
+				handleServerRestart(logMsg);
 			}
 		};
 	}
@@ -178,9 +182,13 @@ export async function createContainerWithAutomaticRestart({
 	// Set up watches
 	function addWatches() {
 		const watcher = restart.container.viteServer.watcher;
-		watcher.on('change', handleServerRestart('Configuration updated. Restarting...'));
-		watcher.on('unlink', handleServerRestart('Configuration removed. Restarting...'));
-		watcher.on('add', handleServerRestart('Configuration added. Restarting...'));
+		watcher.on('change', handleChangeRestart('Configuration updated. Restarting...'));
+		watcher.on('unlink', handleChangeRestart('Configuration removed. Restarting...'));
+		watcher.on('add', handleChangeRestart('Configuration added. Restarting...'));
+
+		// Restart the Astro dev server instead of Vite's when the API is called by plugins.
+		// Ignore the `forceOptimize` parameter for now.
+		restart.container.viteServer.restart = () => handleServerRestart('Restarting...');
 	}
 	addWatches();
 	return restart;
diff --git a/packages/astro/test/units/dev/restart.test.js b/packages/astro/test/units/dev/restart.test.js
index 740ee4976..5c8c4a37c 100644
--- a/packages/astro/test/units/dev/restart.test.js
+++ b/packages/astro/test/units/dev/restart.test.js
@@ -180,4 +180,35 @@ describe('dev container restarts', () => {
 			await restart.container.close();
 		}
 	});
+
+	it('Is able to restart on viteServer.restart API call', async () => {
+		const fs = createFs(
+			{
+				'/src/pages/index.astro': ``,
+			},
+			root
+		);
+
+		const { astroConfig } = await openConfig({
+			cwd: root,
+			flags: {},
+			cmd: 'dev',
+			logging: defaultLogging,
+		});
+		const settings = createSettings(astroConfig, fileURLToPath(root));
+
+		let restart = await createContainerWithAutomaticRestart({
+			params: { fs, root, settings },
+		});
+		await startContainer(restart.container);
+		expect(isStarted(restart.container)).to.equal(true);
+
+		try {
+			let restartComplete = restart.restarted();
+			await restart.container.viteServer.restart();
+			await restartComplete;
+		} finally {
+			await restart.container.close();
+		}
+	});
 });