Improve Astro style HMR for imported styles (#10166)

withastro · Feb 22, 2024 · 598f30c · 598f30c
1 parent 1e638c4
commit 598f30c
Show file tree

Hide file tree

Showing 10 changed files with 141 additions and 108 deletions.
diff --git a/.changeset/little-trains-act.md b/.changeset/little-trains-act.md
@@ -0,0 +1,5 @@
+---
+"astro": patch
+---
+
+Improves Astro style tag HMR when updating imported styles
diff --git a/packages/astro/e2e/fixtures/hmr/src/pages/css-external.astro b/packages/astro/e2e/fixtures/hmr/src/pages/css-external.astro
@@ -0,0 +1,11 @@
+<html>
+	<head>
+		<title>Test</title>
+	</head>
+	<body>
+		<h1 class="css-external">This is blue</h1>
+		<style>
+			@import "../styles/css-external.css";
+		</style>
+	</body>
+</html>
diff --git a/packages/astro/e2e/fixtures/hmr/src/styles/css-external.css b/packages/astro/e2e/fixtures/hmr/src/styles/css-external.css
@@ -0,0 +1,3 @@
+.css-external {
+	color: blue;
+}
diff --git a/packages/astro/e2e/hmr.test.js b/packages/astro/e2e/hmr.test.js
@@ -10,10 +10,18 @@ const test = testFactory({
 
 let devServer;
 
+function throwPageShouldNotReload() {
+	throw new Error('Page should not reload in HMR');
+}
+
 test.beforeAll(async ({ astro }) => {
 	devServer = await astro.startDevServer();
 });
 
+test.afterEach(({ page }) => {
+	page.off('load', throwPageShouldNotReload);
+});
+
 test.afterAll(async () => {
 	await devServer.stop();
 });
@@ -33,15 +41,32 @@ test.describe('Scripts with dependencies', () => {
 	});
 });
 
-test.describe('Styles with dependencies', () => {
-	test('refresh with HMR', async ({ page, astro }) => {
+test.describe('Styles', () => {
+	test('dependencies cause refresh with HMR', async ({ page, astro }) => {
 		await page.goto(astro.resolveUrl('/css-dep'));
 
+		page.once('load', throwPageShouldNotReload);
+
 		const h = page.locator('h1');
 		await expect(h).toHaveCSS('color', 'rgb(0, 0, 255)');
 
 		await astro.editFile('./src/styles/vars.scss', (original) => original.replace('blue', 'red'));
 
 		await expect(h).toHaveCSS('color', 'rgb(255, 0, 0)');
 	});
+
+	test('external CSS refresh with HMR', async ({ page, astro }) => {
+		await page.goto(astro.resolveUrl('/css-external'));
+
+		page.once('load', throwPageShouldNotReload);
+
+		const h = page.locator('h1');
+		await expect(h).toHaveCSS('color', 'rgb(0, 0, 255)');
+
+		await astro.editFile('./src/styles/css-external.css', (original) =>
+			original.replace('blue', 'red')
+		);
+
+		await expect(h).toHaveCSS('color', 'rgb(255, 0, 0)');
+	});
 });
diff --git a/packages/astro/src/core/compile/compile.ts b/packages/astro/src/core/compile/compile.ts
@@ -20,9 +20,16 @@ export interface CompileProps {
 	source: string;
 }
 
-export interface CompileResult extends TransformResult {
-	cssDeps: Set<string>;
-	source: string;
+export interface CompileCssResult {
+	code: string;
+	/**
+	 * The dependencies of the transformed CSS (Normalized paths)
+	 */
+	dependencies?: string[];
+}
+
+export interface CompileResult extends Omit<TransformResult, 'css'> {
+	css: CompileCssResult[];
 }
 
 export async function compile({
@@ -32,7 +39,10 @@ export async function compile({
 	filename,
 	source,
 }: CompileProps): Promise<CompileResult> {
-	const cssDeps = new Set<string>();
+	// Because `@astrojs/compiler` can't return the dependencies for each style transformed,
+	// we need to use an external array to track the dependencies whenever preprocessing is called,
+	// and we'll rebuild the final `css` result after transformation.
+	const cssDeps: CompileCssResult['dependencies'][] = [];
 	const cssTransformErrors: AstroError[] = [];
 	let transformResult: TransformResult;
 
@@ -82,8 +92,10 @@ export async function compile({
 
 	return {
 		...transformResult,
-		cssDeps,
-		source,
+		css: transformResult.css.map((code, i) => ({
+			code,
+			dependencies: cssDeps[i],
+		})),
 	};
 }
 

diff --git a/packages/astro/src/core/compile/style.ts b/packages/astro/src/core/compile/style.ts
@@ -1,7 +1,8 @@
 import type { TransformOptions } from '@astrojs/compiler';
 import fs from 'node:fs';
-import { preprocessCSS, type ResolvedConfig } from 'vite';
+import { normalizePath, preprocessCSS, type ResolvedConfig } from 'vite';
 import { AstroErrorData, CSSError, positionAt } from '../errors/index.js';
+import type { CompileCssResult } from './compile.js';
 
 export function createStylePreprocessor({
 	filename,
@@ -11,18 +12,21 @@ export function createStylePreprocessor({
 }: {
 	filename: string;
 	viteConfig: ResolvedConfig;
-	cssDeps: Set<string>;
+	cssDeps: CompileCssResult['dependencies'][];
 	cssTransformErrors: Error[];
 }): TransformOptions['preprocessStyle'] {
+	let processedStylesCount = 0;
+
 	return async (content, attrs) => {
+		const index = processedStylesCount++;
 		const lang = `.${attrs?.lang || 'css'}`.toLowerCase();
-		const id = `${filename}?astro&type=style&lang${lang}`;
+		const id = `${filename}?astro&type=style&index=${index}&lang${lang}`;
 		try {
 			const result = await preprocessCSS(content, id, viteConfig);
 
-			result.deps?.forEach((dep) => {
-				cssDeps.add(dep);
-			});
+			if (result.deps) {
+				cssDeps[index] = [...result.deps].map((dep) => normalizePath(dep));
+			}
 
 			let map: string | undefined;
 			if (result.map) {

diff --git a/packages/astro/src/vite-plugin-astro/hmr.ts b/packages/astro/src/vite-plugin-astro/hmr.ts
@@ -1,62 +1,45 @@
-import path from 'node:path';
-import { appendForwardSlash } from '@astrojs/internal-helpers/path';
 import type { HmrContext } from 'vite';
 import type { Logger } from '../core/logger/core.js';
-import type { CompileAstroResult } from './compile.js';
-import type { CompileMetadata } from './types.js';
 import { frontmatterRE } from './utils.js';
+import type { CompileMetadata } from './types.js';
 
 export interface HandleHotUpdateOptions {
 	logger: Logger;
-	compile: (code: string, filename: string) => Promise<CompileAstroResult>;
-	astroFileToCssAstroDeps: Map<string, Set<string>>;
 	astroFileToCompileMetadata: Map<string, CompileMetadata>;
 }
 
 export async function handleHotUpdate(
 	ctx: HmrContext,
-	{ logger, compile, astroFileToCssAstroDeps, astroFileToCompileMetadata }: HandleHotUpdateOptions
+	{ logger, astroFileToCompileMetadata }: HandleHotUpdateOptions
 ) {
-	const oldCode = astroFileToCompileMetadata.get(ctx.file)?.originalCode;
-	const newCode = await ctx.read();
+	// HANDLING 1: Invalidate compile metadata if CSS dependency updated
+	//
+	// If any `ctx.file` is part of a CSS dependency of any Astro file, invalidate its `astroFileToCompileMetadata`
+	// so the next transform of the Astro file or Astro script/style virtual module will re-generate it
+	for (const [astroFile, compileData] of astroFileToCompileMetadata) {
+		const isUpdatedFileCssDep = compileData.css.some((css) => css.dependencies?.includes(ctx.file));
+		if (isUpdatedFileCssDep) {
+			astroFileToCompileMetadata.delete(astroFile);
+		}
+	}
+
+	// HANDLING 2: Only invalidate Astro style virtual module if only style tags changed
+	//
 	// If only the style code has changed, e.g. editing the `color`, then we can directly invalidate
 	// the Astro CSS virtual modules only. The main Astro module's JS result will be the same and doesn't
 	// need to be invalidated.
-	if (oldCode && isStyleOnlyChanged(oldCode, newCode)) {
+	const oldCode = astroFileToCompileMetadata.get(ctx.file)?.originalCode;
+	if (oldCode == null) return;
+	const newCode = await ctx.read();
+
+	if (isStyleOnlyChanged(oldCode, newCode)) {
 		logger.debug('watch', 'style-only change');
-		// Re-compile the main Astro component (even though we know its JS result will be the same)
-		// so that `astroFileToCompileMetadata` gets a fresh set of compile metadata to be used
-		// by the virtual modules later in the `load()` hook.
-		await compile(newCode, ctx.file);
+		// Invalidate its `astroFileToCompileMetadata` so that the next transform of Astro style virtual module
+		// will re-generate it
+		astroFileToCompileMetadata.delete(ctx.file);
 		// Only return the Astro styles that have changed!
 		return ctx.modules.filter((mod) => mod.id?.includes('astro&type=style'));
 	}
-
-	// Edge case handling usually caused by Tailwind creating circular dependencies
-	//
-	// TODO: we can also workaround this with better CSS dependency management for Astro files,
-	// so that changes within style tags are scoped to itself. But it'll take a bit of work.
-	// https://github.com/withastro/astro/issues/9370#issuecomment-1850160421
-	for (const [astroFile, cssAstroDeps] of astroFileToCssAstroDeps) {
-		// If the `astroFile` has a CSS dependency on `ctx.file`, there's a good chance this causes a
-		// circular dependency, which Vite doesn't issue a full page reload. Workaround it by forcing a
-		// full page reload ourselves. (Vite bug)
-		// https://github.com/vitejs/vite/pull/15585
-		if (cssAstroDeps.has(ctx.file)) {
-			// Mimic the HMR log as if this file is updated
-			logger.info('watch', getShortName(ctx.file, ctx.server.config.root));
-			// Invalidate the modules of `astroFile` explicitly as Vite may incorrectly soft-invalidate
-			// the parent if the parent actually imported `ctx.file`, but `this.addWatchFile` was also called
-			// on `ctx.file`. Vite should do a hard-invalidation instead. (Vite bug)
-			const parentModules = ctx.server.moduleGraph.getModulesByFile(astroFile);
-			if (parentModules) {
-				for (const mod of parentModules) {
-					ctx.server.moduleGraph.invalidateModule(mod);
-				}
-			}
-			ctx.server.hot.send({ type: 'full-reload', path: '*' });
-		}
-	}
 }
 
 // Disable eslint as we're not sure how to improve this regex yet
@@ -112,7 +95,3 @@ function isArrayEqual(a: any[], b: any[]) {
 	}
 	return true;
 }
-
-function getShortName(file: string, root: string): string {
-	return file.startsWith(appendForwardSlash(root)) ? path.posix.relative(root, file) : file;
-}
diff --git a/packages/astro/src/vite-plugin-astro/index.ts b/packages/astro/src/vite-plugin-astro/index.ts
@@ -9,6 +9,7 @@ import { normalizeFilename } from '../vite-plugin-utils/index.js';
 import { compileAstro, type CompileAstroResult } from './compile.js';
 import { handleHotUpdate } from './hmr.js';
 import { parseAstroRequest } from './query.js';
+import { loadId } from './utils.js';
 export { getAstroMetadata } from './metadata.js';
 export type { AstroPluginMetadata };
 
@@ -24,14 +25,10 @@ export default function astro({ settings, logger }: AstroPluginOptions): vite.Pl
 	const { config } = settings;
 	let server: vite.ViteDevServer | undefined;
 	let compile: (code: string, filename: string) => Promise<CompileAstroResult>;
-	// Tailwind styles could register Astro files as dependencies of other Astro files,
-	// causing circular imports which trips Vite's HMR. This set is passed to `handleHotUpdate`
-	// to force a page reload when these dependency files are updated
-	// NOTE: We need to initialize a map here and in `buildStart` because our unit tests don't
-	// call `buildStart` (test bug)
-	let astroFileToCssAstroDeps = new Map<string, Set<string>>();
 	// Each Astro file has its own compile metadata so that its scripts and styles virtual module
 	// can retrieve their code from here.
+	// NOTE: We need to initialize a map here and in `buildStart` because our unit tests don't
+	// call `buildStart` (test bug)
 	let astroFileToCompileMetadata = new Map<string, CompileMetadata>();
 
 	// Variables for determining if an id starts with /src...
@@ -65,7 +62,6 @@ export default function astro({ settings, logger }: AstroPluginOptions): vite.Pl
 			});
 		},
 		buildStart() {
-			astroFileToCssAstroDeps = new Map();
 			astroFileToCompileMetadata = new Map();
 
 			// Share the `astroFileToCompileMetadata` across the same Astro config as Astro performs
@@ -86,10 +82,19 @@ export default function astro({ settings, logger }: AstroPluginOptions): vite.Pl
 
 			// Astro scripts and styles virtual module code comes from the main Astro compilation
 			// through the metadata from `astroFileToCompileMetadata`. It should always exist as Astro
-			// modules are compiled first, then its virtual modules. If the virtual modules are somehow
-			// compiled first, throw an error and we should investigate it.
+			// modules are compiled first, then its virtual modules.
 			const filename = normalizePath(normalizeFilename(parsedId.filename, config.root));
-			const compileMetadata = astroFileToCompileMetadata.get(filename);
+			let compileMetadata = astroFileToCompileMetadata.get(filename);
+			// If `compileMetadata` doesn't exist in dev, that means the virtual module may have been invalidated.
+			// We try to re-compile the main Astro module (`filename`) first before retrieving the metadata again.
+			if (!compileMetadata && server) {
+				const code = await loadId(server.pluginContainer, filename);
+				// `compile` should re-set `filename` in `astroFileToCompileMetadata`
+				if (code != null) await compile(code, filename);
+				compileMetadata = astroFileToCompileMetadata.get(filename);
+			}
+			// If the metadata still doesn't exist, that means the virtual modules are somehow compiled first,
+			// throw an error and we should investigate it.
 			if (!compileMetadata) {
 				throw new Error(
 					`No cached compile metadata found for "${id}". The main Astro module "${filename}" should have ` +
@@ -103,12 +108,15 @@ export default function astro({ settings, logger }: AstroPluginOptions): vite.Pl
 						throw new Error(`Requests for Astro CSS must include an index.`);
 					}
 
-					const code = compileMetadata.css[query.index];
-					if (!code) {
+					const result = compileMetadata.css[query.index];
+					if (!result) {
 						throw new Error(`No Astro CSS at index ${query.index}`);
 					}
 
-					return { code };
+					// Register dependencies from preprocessing this style
+					result.dependencies?.forEach((dep) => this.addWatchFile(dep));
+
+					return { code: result.code };
 				}
 				case 'script': {
 					if (typeof query.index === 'undefined') {
@@ -174,34 +182,6 @@ export default function astro({ settings, logger }: AstroPluginOptions): vite.Pl
 			const filename = normalizePath(parsedId.filename);
 			const transformResult = await compile(source, filename);
 
-			// Register dependencies of this module
-			const astroDeps = new Set<string>();
-			for (const dep of transformResult.cssDeps) {
-				if (dep.endsWith('.astro')) {
-					astroDeps.add(dep);
-				}
-				this.addWatchFile(dep);
-			}
-			astroFileToCssAstroDeps.set(id, astroDeps);
-
-			// When a dependency from the styles are updated, the dep and Astro module will get invalidated.
-			// However, the Astro style virtual module is not invalidated because we didn't register that the virtual
-			// module has that dependency. We currently can't do that either because of a Vite bug.
-			// https://github.com/vitejs/vite/pull/15608
-			// Here we manually invalidate the virtual modules ourselves when we're compiling the Astro module.
-			// When that bug is resolved, we can add the dependencies to the virtual module directly and remove this.
-			if (server) {
-				const mods = server.moduleGraph.getModulesByFile(filename);
-				if (mods) {
-					const seen = new Set(mods);
-					for (const mod of mods) {
-						if (mod.url.includes('?astro')) {
-							server.moduleGraph.invalidateModule(mod, seen);
-						}
-					}
-				}
-			}
-
 			const astroMetadata: AstroPluginMetadata['astro'] = {
 				clientOnlyComponents: transformResult.clientOnlyComponents,
 				hydratedComponents: transformResult.hydratedComponents,
@@ -225,14 +205,7 @@ export default function astro({ settings, logger }: AstroPluginOptions): vite.Pl
 			};
 		},
 		async handleHotUpdate(ctx) {
-			if (!ctx.file.endsWith('.astro')) return;
-
-			return handleHotUpdate(ctx, {
-				logger,
-				compile,
-				astroFileToCssAstroDeps,
-				astroFileToCompileMetadata,
-			});
+			return handleHotUpdate(ctx, { logger, astroFileToCompileMetadata });
 		},
 	};
 

diff --git a/packages/astro/src/vite-plugin-astro/types.ts b/packages/astro/src/vite-plugin-astro/types.ts
@@ -1,5 +1,6 @@
 import type { HoistedScript, TransformResult } from '@astrojs/compiler';
 import type { PropagationHint } from '../@types/astro.js';
+import type { CompileCssResult } from '../core/compile/compile.js';
 
 export interface PageOptions {
 	prerender?: boolean;
@@ -20,7 +21,7 @@ export interface CompileMetadata {
 	/** Used for HMR to compare code changes */
 	originalCode: string;
 	/** For Astro CSS virtual module */
-	css: string[];
+	css: CompileCssResult[];
 	/** For Astro hoisted scripts virtual module */
 	scripts: HoistedScript[];
 }