From 27ee8b97ae3868d93bcd6c9302a401901394018f Mon Sep 17 00:00:00 2001
From: Nate Moore <natemoo-re@users.noreply.github.com>
Date: Fri, 22 Jul 2022 13:37:20 -0500
Subject: [PATCH] Overhaul HMR handling for `.astro` files (#3932)

* feat(hmr): overhaul HMR handling for `.astro` files

* test: remove irrelevant hmr-css test

* fix: ignore HMR logs for monorepo files

* fix: explicitly call import.meta.hot.decline()

* chore: add changeset

* fix: run smoke tests in serial

* chore: bump postcss test timeout

* test: fix undefined `this`

Co-authored-by: Nate Moore <nate@astro.build>
---
 .changeset/popular-taxis-prove.md             |   5 +
 package.json                                  |   2 +-
 packages/astro/src/core/render/dev/hmr.ts     |  11 --
 packages/astro/src/core/render/dev/index.ts   |  14 ++-
 packages/astro/src/runtime/client/hmr.ts      | 110 +-----------------
 .../src/vite-plugin-astro-server/index.ts     |  29 -----
 packages/astro/src/vite-plugin-astro/hmr.ts   |  17 ++-
 packages/astro/src/vite-plugin-astro/index.ts |   9 +-
 .../fixtures/hmr-css/src/pages/index.astro    |  11 --
 packages/astro/test/hmr-css.test.js           |  34 ------
 packages/astro/test/postcss.test.js           |   3 +-
 11 files changed, 41 insertions(+), 204 deletions(-)
 create mode 100644 .changeset/popular-taxis-prove.md
 delete mode 100644 packages/astro/src/core/render/dev/hmr.ts
 delete mode 100644 packages/astro/test/fixtures/hmr-css/src/pages/index.astro
 delete mode 100644 packages/astro/test/hmr-css.test.js

diff --git a/.changeset/popular-taxis-prove.md b/.changeset/popular-taxis-prove.md
new file mode 100644
index 000000000000..101dfd7ddaa7
--- /dev/null
+++ b/.changeset/popular-taxis-prove.md
@@ -0,0 +1,5 @@
+---
+'astro': patch
+---
+
+Overhaul HMR handling for more stable live reload behavior
diff --git a/package.json b/package.json
index 0595ac436c80..55263594adc4 100644
--- a/package.json
+++ b/package.json
@@ -19,7 +19,7 @@
     "test": "turbo run test --output-logs=new-only --concurrency=1",
     "test:match": "cd packages/astro && pnpm run test:match",
     "test:templates": "turbo run test --filter=create-astro --concurrency=1",
-    "test:smoke": "turbo run build --filter=\"@example/*\" --filter=\"astro.build\" --filter=\"docs\" --output-logs=new-only",
+    "test:smoke": "turbo run build --filter=\"@example/*\" --filter=\"astro.build\" --filter=\"docs\" --output-logs=new-only --concurrency=1",
     "test:vite-ci": "turbo run test --output-logs=new-only --no-deps --scope=astro --concurrency=1",
     "test:e2e": "cd packages/astro && pnpm playwright install && pnpm run test:e2e",
     "test:e2e:match": "cd packages/astro && pnpm playwright install && pnpm run test:e2e:match",
diff --git a/packages/astro/src/core/render/dev/hmr.ts b/packages/astro/src/core/render/dev/hmr.ts
deleted file mode 100644
index 3c795fdb139a..000000000000
--- a/packages/astro/src/core/render/dev/hmr.ts
+++ /dev/null
@@ -1,11 +0,0 @@
-import fs from 'fs';
-import { fileURLToPath } from 'url';
-
-let hmrScript: string;
-export async function getHmrScript() {
-	if (hmrScript) return hmrScript;
-	const filePath = fileURLToPath(new URL('../../../runtime/client/hmr.js', import.meta.url));
-	const content = await fs.promises.readFile(filePath);
-	hmrScript = content.toString();
-	return hmrScript;
-}
diff --git a/packages/astro/src/core/render/dev/index.ts b/packages/astro/src/core/render/dev/index.ts
index e3b6f0ac7c96..428c30edf10c 100644
--- a/packages/astro/src/core/render/dev/index.ts
+++ b/packages/astro/src/core/render/dev/index.ts
@@ -150,11 +150,19 @@ export async function render(
 
 	let styles = new Set<SSRElement>();
 	[...stylesMap].forEach(([url, content]) => {
-		// The URL is only used by HMR for Svelte components
-		// See src/runtime/client/hmr.ts for more details
+		// Vite handles HMR for styles injected as scripts
+		scripts.add({
+			props: {
+				type: 'module',
+				src: url,
+				'data-astro-injected': true,
+			},
+			children: '',
+		});
+		// But we still want to inject the styles to avoid FOUC
 		styles.add({
 			props: {
-				'data-astro-injected': svelteStylesRE.test(url) ? url : true,
+				'data-astro-injected': url,
 			},
 			children: content,
 		});
diff --git a/packages/astro/src/runtime/client/hmr.ts b/packages/astro/src/runtime/client/hmr.ts
index 98cf839afcba..98153f4b262e 100644
--- a/packages/astro/src/runtime/client/hmr.ts
+++ b/packages/astro/src/runtime/client/hmr.ts
@@ -1,59 +1,8 @@
 /// <reference types="vite/client" />
 if (import.meta.hot) {
-	import.meta.hot.accept((mod) => mod);
-
-	const parser = new DOMParser();
-
-	const KNOWN_MANUAL_HMR_EXTENSIONS = new Set(['.astro', '.md', '.mdx']);
-	function needsManualHMR(path: string) {
-		for (const ext of KNOWN_MANUAL_HMR_EXTENSIONS.values()) {
-			if (path.endsWith(ext)) return true;
-		}
-		return false;
-	}
-
-	async function updatePage() {
-		const { default: diff } = await import('micromorph');
-		const html = await fetch(`${window.location}`).then((res) => res.text());
-		const doc = parser.parseFromString(html, 'text/html');
-		for (const style of sheetsMap.values()) {
-			doc.head.appendChild(style);
-		}
-		// Match incoming islands to current state
-		for (const root of doc.querySelectorAll('astro-island')) {
-			const uid = root.getAttribute('uid');
-			const current = document.querySelector(`astro-island[uid="${uid}"]`);
-			if (current) {
-				current.setAttribute('data-persist', '');
-				root.replaceWith(current);
-			}
-		}
-		// both Vite and Astro's HMR scripts include `type="text/css"` on injected
-		// <style> blocks. These style blocks would not have been rendered in Astro's
-		// build and need to be persisted when diffing HTML changes.
-		for (const style of document.querySelectorAll("style[type='text/css']")) {
-			style.setAttribute('data-persist', '');
-			doc.head.appendChild(style.cloneNode(true));
-		}
-		return diff(document, doc).then(() => {
-			// clean up data-persist attributes added before diffing
-			for (const root of document.querySelectorAll('astro-island[data-persist]')) {
-				root.removeAttribute('data-persist');
-			}
-			for (const style of document.querySelectorAll("style[type='text/css'][data-persist]")) {
-				style.removeAttribute('data-persist');
-			}
-		});
-	}
-	async function updateAll(files: any[]) {
-		let hasManualUpdate = false;
-		let styles = [];
-		for (const file of files) {
-			if (needsManualHMR(file.acceptedPath)) {
-				hasManualUpdate = true;
-				continue;
-			}
-			if (file.acceptedPath.includes('svelte&type=style')) {
+	import.meta.hot.on('vite:beforeUpdate', async (payload) => {
+		for (const file of payload.updates) {
+			if (file.acceptedPath.includes('svelte&type=style') || file.acceptedPath.includes('astro&type=style')) {
 				// This will only be called after the svelte component has hydrated in the browser.
 				// At this point Vite is tracking component style updates, we need to remove
 				// styles injected by Astro for the component in favor of Vite's internal HMR.
@@ -70,59 +19,6 @@ if (import.meta.hot) {
 					link.replaceWith(link.cloneNode(true));
 				}
 			}
-			if (file.acceptedPath.includes('astro&type=style')) {
-				styles.push(
-					fetch(file.acceptedPath)
-						.then((res) => res.text())
-						.then((res) => [file.acceptedPath, res])
-				);
-			}
-		}
-		if (styles.length > 0) {
-			for (const [id, content] of await Promise.all(styles)) {
-				updateStyle(id, content);
-			}
 		}
-		if (hasManualUpdate) {
-			return await updatePage();
-		}
-	}
-	import.meta.hot.on('vite:beforeUpdate', async (event) => {
-		await updateAll(event.updates);
 	});
 }
-
-const sheetsMap = new Map();
-
-function updateStyle(id: string, content: string): void {
-	let style = sheetsMap.get(id);
-	if (style && !(style instanceof HTMLStyleElement)) {
-		removeStyle(id);
-		style = undefined;
-	}
-
-	if (!style) {
-		style = document.createElement('style');
-		style.setAttribute('type', 'text/css');
-		style.innerHTML = content;
-		document.head.appendChild(style);
-	} else {
-		style.innerHTML = content;
-	}
-	sheetsMap.set(id, style);
-}
-
-function removeStyle(id: string): void {
-	const style = sheetsMap.get(id);
-	if (style) {
-		if (style instanceof CSSStyleSheet) {
-			// @ts-expect-error: using experimental API
-			document.adoptedStyleSheets = document.adoptedStyleSheets.filter(
-				(s: CSSStyleSheet) => s !== style
-			);
-		} else {
-			document.head.removeChild(style);
-		}
-		sheetsMap.delete(id);
-	}
-}
diff --git a/packages/astro/src/vite-plugin-astro-server/index.ts b/packages/astro/src/vite-plugin-astro-server/index.ts
index 84a92f00de1a..86012991ef09 100644
--- a/packages/astro/src/vite-plugin-astro-server/index.ts
+++ b/packages/astro/src/vite-plugin-astro-server/index.ts
@@ -350,31 +350,6 @@ async function handleRequest(
 	}
 }
 
-/**
- * Vite HMR sends requests for new CSS and those get returned as JS, but we want it to be CSS
- * since they are inside of a link tag for Astro.
- */
-const forceTextCSSForStylesMiddleware: vite.Connect.NextHandleFunction = function (req, res, next) {
-	if (req.url) {
-		// We are just using this to parse the URL to get the search params object
-		// so the second arg here doesn't matter
-		const url = new URL(req.url, 'https://astro.build');
-		// lang.css is a search param that exists on Astro, Svelte, and Vue components.
-		// We only want to override for astro files.
-		if (url.searchParams.has('astro') && url.searchParams.has('lang.css')) {
-			// Override setHeader so we can set the correct content-type for this request.
-			const setHeader = res.setHeader;
-			res.setHeader = function (key, value) {
-				if (key.toLowerCase() === 'content-type') {
-					return setHeader.call(this, key, 'text/css');
-				}
-				return setHeader.apply(this, [key, value]);
-			};
-		}
-	}
-	next();
-};
-
 export default function createPlugin({ config, logging }: AstroPluginOptions): vite.Plugin {
 	return {
 		name: 'astro:server',
@@ -396,10 +371,6 @@ export default function createPlugin({ config, logging }: AstroPluginOptions): v
 				removeViteHttpMiddleware(viteServer.middlewares);
 
 				// Push this middleware to the front of the stack so that it can intercept responses.
-				viteServer.middlewares.stack.unshift({
-					route: '',
-					handle: forceTextCSSForStylesMiddleware,
-				});
 				if (config.base !== '/') {
 					viteServer.middlewares.stack.unshift({
 						route: '',
diff --git a/packages/astro/src/vite-plugin-astro/hmr.ts b/packages/astro/src/vite-plugin-astro/hmr.ts
index 3d990fe4bbce..e06325bc33d0 100644
--- a/packages/astro/src/vite-plugin-astro/hmr.ts
+++ b/packages/astro/src/vite-plugin-astro/hmr.ts
@@ -2,6 +2,7 @@ import type { PluginContext as RollupPluginContext, ResolvedId } from 'rollup';
 import type { HmrContext, ModuleNode, ViteDevServer } from 'vite';
 import type { AstroConfig } from '../@types/astro';
 import type { LogOptions } from '../core/logger/core.js';
+import { fileURLToPath } from 'node:url';
 import { info } from '../core/logger/core.js';
 import * as msg from '../core/messages.js';
 import { invalidateCompilation, isCached } from './compile.js';
@@ -49,21 +50,31 @@ export async function trackCSSDependencies(
 	}
 }
 
+const PKG_PREFIX = new URL('../../', import.meta.url)
+const isPkgFile = (id: string|null) => {
+	return id?.startsWith(fileURLToPath(PKG_PREFIX)) || id?.startsWith(PKG_PREFIX.pathname)
+}
+
 export async function handleHotUpdate(ctx: HmrContext, config: AstroConfig, logging: LogOptions) {
 	// Invalidate the compilation cache so it recompiles
 	invalidateCompilation(config, ctx.file);
+	
+	// Skip monorepo files to avoid console spam
+	if (isPkgFile(ctx.file)) {
+		return;
+	}
 
 	// go through each of these modules importers and invalidate any .astro compilation
 	// that needs to be rerun.
 	const filtered = new Set<ModuleNode>(ctx.modules);
 	const files = new Set<string>();
 	for (const mod of ctx.modules) {
-		// This is always the HMR script, we skip it to avoid spamming
-		// the browser console with HMR updates about this file
-		if (mod.id?.endsWith('.astro?html-proxy&index=0.js')) {
+		// Skip monorepo files to avoid console spam
+		if (isPkgFile(mod.id ?? mod.file)) {
 			filtered.delete(mod);
 			continue;
 		}
+
 		if (mod.file && isCached(config, mod.file)) {
 			filtered.add(mod);
 			files.add(mod.file);
diff --git a/packages/astro/src/vite-plugin-astro/index.ts b/packages/astro/src/vite-plugin-astro/index.ts
index f47940c0c22b..b4925d0fd66e 100644
--- a/packages/astro/src/vite-plugin-astro/index.ts
+++ b/packages/astro/src/vite-plugin-astro/index.ts
@@ -268,16 +268,17 @@ export default function astro({ config, logging }: AstroPluginOptions): vite.Plu
 						SUFFIX += `import "${id}?astro&type=script&index=${i}&lang.ts";`;
 						i++;
 					}
-
-					SUFFIX += `\nif (import.meta.hot) {
-						import.meta.hot.accept(mod => mod);
-					}`;
 				}
 				// Add handling to inject scripts into each page JS bundle, if needed.
 				if (isPage) {
 					SUFFIX += `\nimport "${PAGE_SSR_SCRIPT_ID}";`;
 				}
 
+				// Prefer live reload to HMR in `.astro` files
+				if (!resolvedConfig.isProduction) {
+					SUFFIX += `\nif (import.meta.hot) { import.meta.hot.decline() }`;
+				}
+
 				const astroMetadata: AstroPluginMetadata['astro'] = {
 					clientOnlyComponents: transformResult.clientOnlyComponents,
 					hydratedComponents: transformResult.hydratedComponents,
diff --git a/packages/astro/test/fixtures/hmr-css/src/pages/index.astro b/packages/astro/test/fixtures/hmr-css/src/pages/index.astro
deleted file mode 100644
index 840e60e01dec..000000000000
--- a/packages/astro/test/fixtures/hmr-css/src/pages/index.astro
+++ /dev/null
@@ -1,11 +0,0 @@
-<html>
-	<head>
-		<title>Testing</title>
-		<style>
-			background { background: brown; }
-		</style>
-	</head>
-	<body>
-		<h1>Testing</h1>
-	</body>
-</html>
diff --git a/packages/astro/test/hmr-css.test.js b/packages/astro/test/hmr-css.test.js
deleted file mode 100644
index b2b4341b4d00..000000000000
--- a/packages/astro/test/hmr-css.test.js
+++ /dev/null
@@ -1,34 +0,0 @@
-import { isWindows, loadFixture } from './test-utils.js';
-import { expect } from 'chai';
-import * as cheerio from 'cheerio';
-
-describe('HMR - CSS', () => {
-	if (isWindows) return;
-
-	/** @type {import('./test-utils').Fixture} */
-	let fixture;
-	/** @type {import('./test-utils').DevServer} */
-	let devServer;
-
-	before(async () => {
-		fixture = await loadFixture({
-			root: './fixtures/hmr-css/',
-		});
-		devServer = await fixture.startDevServer();
-	});
-
-	after(async () => {
-		await devServer.stop();
-	});
-
-	it('Timestamp URL used by Vite gets the right mime type', async () => {
-		// Index page is always loaded first by the browser
-		await fixture.fetch('/');
-		// Now we can simulate what happens in the browser
-		let res = await fixture.fetch(
-			'/src/pages/index.astro?astro=&type=style&index=0&lang.css=&t=1653657441095'
-		);
-		let headers = res.headers;
-		expect(headers.get('content-type')).to.equal('text/css');
-	});
-});
diff --git a/packages/astro/test/postcss.test.js b/packages/astro/test/postcss.test.js
index 1cf06bee1029..28de600da4a5 100644
--- a/packages/astro/test/postcss.test.js
+++ b/packages/astro/test/postcss.test.js
@@ -3,12 +3,13 @@ import * as cheerio from 'cheerio';
 import eol from 'eol';
 import { loadFixture } from './test-utils.js';
 
-describe('PostCSS', () => {
+describe('PostCSS', function () {
 	const PREFIXED_CSS = `{-webkit-appearance:none;appearance:none`;
 
 	let fixture;
 	let bundledCSS;
 	before(async () => {
+		this.timeout(45000); // test needs a little more time in CI
 		fixture = await loadFixture({
 			root: './fixtures/postcss',
 		});