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

withastro · Jul 19, 2022 · e4c574c · e4c574c
1 parent 4412fe6
commit e4c574c
Show file tree

Hide file tree

Showing 6 changed files with 15 additions and 156 deletions.
diff --git a/packages/astro/src/core/render/dev/hmr.ts b/packages/astro/src/core/render/dev/hmr.ts
diff --git 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
@@ -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
@@ -331,31 +331,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',
@@ -376,11 +351,6 @@ export default function createPlugin({ config, logging }: AstroPluginOptions): v
 			return () => {
 				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,
-				});
 				viteServer.middlewares.use(async (req, res) => {
 					if (!req.url || !req.method) {
 						throw new Error('Incomplete request');

diff --git a/packages/astro/src/vite-plugin-astro/hmr.ts b/packages/astro/src/vite-plugin-astro/hmr.ts
@@ -60,7 +60,7 @@ export async function handleHotUpdate(ctx: HmrContext, config: AstroConfig, logg
 	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')) {
+		if (mod.id?.endsWith('/astro/dist/runtime/client/hmr.js')) {
 			filtered.delete(mod);
 			continue;
 		}

diff --git a/packages/astro/src/vite-plugin-astro/index.ts b/packages/astro/src/vite-plugin-astro/index.ts
@@ -256,10 +256,6 @@ 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) {