Refactor CSS preprocess and HMR (#4422)

6 files changed, 118 insertions, 142 deletions

diff --git a/.changeset/tricky-melons-grin.md b/.changeset/tricky-melons-grin.md
new file mode 100644
index 000000000..ea85c3463
--- /dev/null
+++ b/.changeset/tricky-melons-grin.md
@@ -0,0 +1,5 @@
+---
+'astro': patch
+---
+
+Refactor CSS preprocess and deps HMR
diff --git a/packages/astro/src/vite-plugin-astro/compile.ts b/packages/astro/src/vite-plugin-astro/compile.ts
index 02c3753c4..368e310fc 100644
--- a/packages/astro/src/vite-plugin-astro/compile.ts
+++ b/packages/astro/src/vite-plugin-astro/compile.ts
@@ -1,40 +1,21 @@
 import type { TransformResult } from '@astrojs/compiler';
 import type { PluginContext, SourceMapInput } from 'rollup';
+import type { ViteDevServer } from 'vite';
 import type { AstroConfig } from '../@types/astro';
-import type { TransformHook } from './styles';
+import type { TransformStyleWithVite } from './styles';
 
 import { transform } from '@astrojs/compiler';
 import { fileURLToPath } from 'url';
 import { AstroErrorCodes } from '../core/errors.js';
 import { prependForwardSlash } from '../core/path.js';
 import { viteID } from '../core/util.js';
-import { transformWithVite } from './styles.js';
 
 type CompilationCache = Map<string, CompileResult>;
 type CompileResult = TransformResult & {
-	rawCSSDeps: Set<string>;
+	cssDeps: Set<string>;
 	source: string;
 };
 
-/**
- * Note: this is currently needed because Astro is directly using a Vite internal CSS transform. This gives us
- * some nice features out of the box, but at the expense of also running Vite's CSS postprocessing build step,
- * which does some things that we don't like, like resolving/handling `@import` too early. This function pulls
- * out the `@import` tags to be added back later, and then finally handled correctly by Vite.
- *
- * In the future, we should remove this workaround and most likely implement our own Astro style handling without
- * having to hook into Vite's internals.
- */
-function createImportPlaceholder(spec: string) {
-	// Note: We keep this small so that we can attempt to exactly match the # of characters in the original @import.
-	// This keeps sourcemaps accurate (to the best of our ability) at the intermediate step where this appears.
-	// ->  `@import '${spec}';`;
-	return `/*IMPORT:${spec}*/`;
-}
-function safelyReplaceImportPlaceholder(code: string) {
-	return code.replace(/\/\*IMPORT\:(.*?)\*\//g, `@import '$1';`);
-}
-
 const configCache = new WeakMap<AstroConfig, CompilationCache>();
 
 export interface CompileProps {
@@ -43,7 +24,8 @@ export interface CompileProps {
 	moduleId: string;
 	source: string;
 	ssr: boolean;
-	viteTransform: TransformHook;
+	transformStyleWithVite: TransformStyleWithVite;
+	viteDevServer?: ViteDevServer;
 	pluginContext: PluginContext;
 }
 
@@ -63,13 +45,19 @@ async function compile({
 	moduleId,
 	source,
 	ssr,
-	viteTransform,
+	transformStyleWithVite,
+	viteDevServer,
 	pluginContext,
 }: CompileProps): Promise<CompileResult> {
 	const normalizedID = getNormalizedID(filename);
-	let rawCSSDeps = new Set<string>();
+	let cssDeps = new Set<string>();
 	let cssTransformError: Error | undefined;
 
+	// handleHotUpdate doesn't have `addWatchFile` used by transformStyleWithVite.
+	if (!pluginContext.addWatchFile) {
+		pluginContext.addWatchFile = () => {};
+	}
+
 	// Transform from `.astro` to valid `.ts`
 	// use `sourcemap: "both"` so that sourcemap is included in the code
 	// result passed to esbuild, but also available in the catch handler.
@@ -89,32 +77,21 @@ async function compile({
 			const lang = `.${attrs?.lang || 'css'}`.toLowerCase();
 
 			try {
-				// In the static build, grab any @import as CSS dependencies for HMR.
-				value.replace(
-					/(?:@import)\s(?:url\()?\s?["\'](.*?)["\']\s?\)?(?:[^;]*);?/gi,
-					(match, spec) => {
-						rawCSSDeps.add(spec);
-						// If the language is CSS: prevent `@import` inlining to prevent scoping of imports.
-						// Otherwise: Sass, etc. need to see imports for variables, so leave in for their compiler to handle.
-						if (lang === '.css') {
-							return createImportPlaceholder(spec);
-						} else {
-							return match;
-						}
-					}
-				);
-
-				const result = await transformWithVite({
-					value,
-					lang,
+				const result = await transformStyleWithVite.call(pluginContext, {
 					id: normalizedID,
-					transformHook: viteTransform,
+					source: value,
+					lang,
 					ssr,
-					pluginContext,
+					viteDevServer,
 				});
 
-				let map: SourceMapInput | undefined;
 				if (!result) return null as any; // TODO: add type in compiler to fix "any"
+
+				for (const dep of result.deps) {
+					cssDeps.add(dep);
+				}
+
+				let map: SourceMapInput | undefined;
 				if (result.map) {
 					if (typeof result.map === 'string') {
 						map = result.map;
@@ -122,8 +99,8 @@ async function compile({
 						map = result.map.toString();
 					}
 				}
-				const code = safelyReplaceImportPlaceholder(result.code);
-				return { code, map };
+
+				return { code: result.code, map };
 			} catch (err) {
 				// save error to throw in plugin context
 				cssTransformError = err as any;
@@ -147,8 +124,8 @@ async function compile({
 		});
 
 	const compileResult: CompileResult = Object.create(transformResult, {
-		rawCSSDeps: {
-			value: rawCSSDeps,
+		cssDeps: {
+			value: cssDeps,
 		},
 		source: {
 			value: source,
diff --git a/packages/astro/src/vite-plugin-astro/hmr.ts b/packages/astro/src/vite-plugin-astro/hmr.ts
index c218998fa..6ee123e07 100644
--- a/packages/astro/src/vite-plugin-astro/hmr.ts
+++ b/packages/astro/src/vite-plugin-astro/hmr.ts
@@ -1,6 +1,5 @@
 import { fileURLToPath } from 'node:url';
-import type { PluginContext as RollupPluginContext, ResolvedId } from 'rollup';
-import type { HmrContext, ModuleNode, ViteDevServer } from 'vite';
+import type { HmrContext, ModuleNode } from 'vite';
 import type { AstroConfig } from '../@types/astro';
 import type { LogOptions } from '../core/logger/core.js';
 import { info } from '../core/logger/core.js';
@@ -8,49 +7,6 @@ import * as msg from '../core/messages.js';
 import { cachedCompilation, invalidateCompilation, isCached } from './compile.js';
 import { isAstroScript } from './query.js';
 
-interface TrackCSSDependenciesOptions {
-	viteDevServer: ViteDevServer | null;
-	filename: string;
-	id: string;
-	deps: Set<string>;
-}
-
-export async function trackCSSDependencies(
-	this: RollupPluginContext,
-	opts: TrackCSSDependenciesOptions
-): Promise<void> {
-	const { viteDevServer, filename, deps, id } = opts;
-	// Dev, register CSS dependencies for HMR.
-	if (viteDevServer) {
-		const mod = viteDevServer.moduleGraph.getModuleById(id);
-		if (mod) {
-			const cssDeps = (
-				await Promise.all(
-					Array.from(deps).map((spec) => {
-						return this.resolve(spec, id);
-					})
-				)
-			)
-				.filter(Boolean)
-				.map((dep) => (dep as ResolvedId).id);
-
-			const { moduleGraph } = viteDevServer;
-			// record deps in the module graph so edits to @import css can trigger
-			// main import to hot update
-			const depModules = new Set(mod.importedModules);
-			for (const dep of cssDeps) {
-				depModules.add(moduleGraph.createFileOnlyEntry(dep));
-			}
-
-			// Update the module graph, telling it about our CSS deps.
-			moduleGraph.updateModuleInfo(mod, depModules, new Map(), new Set(), new Set(), true);
-			for (const dep of cssDeps) {
-				this.addWatchFile(dep);
-			}
-		}
-	}
-}
-
 const PKG_PREFIX = new URL('../../', import.meta.url);
 const isPkgFile = (id: string | null) => {
 	return id?.startsWith(fileURLToPath(PKG_PREFIX)) || id?.startsWith(PKG_PREFIX.pathname);
@@ -125,6 +81,13 @@ export async function handleHotUpdate(
 	for (const file of files) {
 		if (isStyleOnlyChange && file === ctx.file) continue;
 		invalidateCompilation(config, file);
+		// If `ctx.file` is depended by an .astro file, e.g via `this.addWatchFile`,
+		// Vite doesn't trigger updating that .astro file by default. See:
+		// https://github.com/vitejs/vite/issues/3216
+		// For now, we trigger the change manually here.
+		if (file.endsWith('.astro')) {
+			ctx.server.moduleGraph.onFileChange(file);
+		}
 	}
 
 	// Bugfix: sometimes style URLs get normalized and end with `lang.css=`
diff --git a/packages/astro/src/vite-plugin-astro/index.ts b/packages/astro/src/vite-plugin-astro/index.ts
index 393128803..f7a1f527d 100644
--- a/packages/astro/src/vite-plugin-astro/index.ts
+++ b/packages/astro/src/vite-plugin-astro/index.ts
@@ -11,9 +11,9 @@ import { fileURLToPath } from 'url';
 import { isRelativePath, startsWithForwardSlash } from '../core/path.js';
 import { getFileInfo } from '../vite-plugin-utils/index.js';
 import { cachedCompilation, CompileProps, getCachedSource } from './compile.js';
-import { handleHotUpdate, trackCSSDependencies } from './hmr.js';
+import { handleHotUpdate } from './hmr.js';
 import { parseAstroRequest, ParsedRequestResult } from './query.js';
-import { getViteTransform, TransformHook } from './styles.js';
+import { createTransformStyleWithViteFn, TransformStyleWithVite } from './styles.js';
 
 const FRONTMATTER_PARSE_REGEXP = /^\-\-\-(.*)^\-\-\-/ms;
 interface AstroPluginOptions {
@@ -38,8 +38,8 @@ export default function astro({ config, logging }: AstroPluginOptions): vite.Plu
 	}
 
 	let resolvedConfig: vite.ResolvedConfig;
-	let viteTransform: TransformHook;
-	let viteDevServer: vite.ViteDevServer | null = null;
+	let transformStyleWithVite: TransformStyleWithVite;
+	let viteDevServer: vite.ViteDevServer | undefined;
 
 	// Variables for determing if an id starts with /src...
 	const srcRootWeb = config.srcDir.pathname.slice(config.root.pathname.length - 1);
@@ -60,7 +60,7 @@ export default function astro({ config, logging }: AstroPluginOptions): vite.Plu
 		enforce: 'pre', // run transforms before other plugins can
 		configResolved(_resolvedConfig) {
 			resolvedConfig = _resolvedConfig;
-			viteTransform = getViteTransform(resolvedConfig);
+			transformStyleWithVite = createTransformStyleWithViteFn(_resolvedConfig);
 		},
 		configureServer(server) {
 			viteDevServer = server;
@@ -118,7 +118,8 @@ export default function astro({ config, logging }: AstroPluginOptions): vite.Plu
 				moduleId: id,
 				source,
 				ssr: Boolean(opts?.ssr),
-				viteTransform,
+				transformStyleWithVite,
+				viteDevServer,
 				pluginContext: this,
 			};
 
@@ -129,14 +130,6 @@ export default function astro({ config, logging }: AstroPluginOptions): vite.Plu
 					}
 
 					const transformResult = await cachedCompilation(compileProps);
-
-					// Track any CSS dependencies so that HMR is triggered when they change.
-					await trackCSSDependencies.call(this, {
-						viteDevServer,
-						id,
-						filename,
-						deps: transformResult.rawCSSDeps,
-					});
 					const csses = transformResult.css;
 					const code = csses[query.index];
 
@@ -224,7 +217,8 @@ export default function astro({ config, logging }: AstroPluginOptions): vite.Plu
 				moduleId: id,
 				source,
 				ssr: Boolean(opts?.ssr),
-				viteTransform,
+				transformStyleWithVite,
+				viteDevServer,
 				pluginContext: this,
 			};
 
@@ -232,6 +226,10 @@ export default function astro({ config, logging }: AstroPluginOptions): vite.Plu
 				const transformResult = await cachedCompilation(compileProps);
 				const { fileId: file, fileUrl: url } = getFileInfo(id, config);
 
+				for (const dep of transformResult.cssDeps) {
+					this.addWatchFile(dep);
+				}
+
 				// Compile all TypeScript to JavaScript.
 				// Also, catches invalid JS/TS in the compiled output before returning.
 				const { code, map } = await esbuild.transform(transformResult.code, {
@@ -355,7 +353,8 @@ ${source}
 				moduleId: context.file,
 				source: await context.read(),
 				ssr: true,
-				viteTransform,
+				transformStyleWithVite,
+				viteDevServer,
 				pluginContext: this,
 			};
 			const compile = () => cachedCompilation(compileProps);
diff --git a/packages/astro/src/vite-plugin-astro/styles.ts b/packages/astro/src/vite-plugin-astro/styles.ts
index 220929d24..e473531fe 100644
--- a/packages/astro/src/vite-plugin-astro/styles.ts
+++ b/packages/astro/src/vite-plugin-astro/styles.ts
@@ -9,34 +9,61 @@ export type TransformHook = (
 	ssr?: boolean
 ) => Promise<vite.TransformResult>;
 
-/** Load vite:css’ transform() hook */
-export function getViteTransform(viteConfig: vite.ResolvedConfig): TransformHook {
-	const viteCSSPlugin = viteConfig.plugins.find(({ name }) => name === 'vite:css');
-	if (!viteCSSPlugin) throw new Error(`vite:css plugin couldn’t be found`);
-	if (!viteCSSPlugin.transform) throw new Error(`vite:css has no transform() hook`);
-	return viteCSSPlugin.transform as any;
-}
-
-interface TransformWithViteOptions {
-	value: string;
-	lang: string;
+interface TransformStyleWithViteOptions {
 	id: string;
-	transformHook: TransformHook;
-	pluginContext: PluginContext;
+	source: string;
+	lang: string;
 	ssr?: boolean;
+	viteDevServer?: vite.ViteDevServer;
+}
+
+export interface TransformStyleWithVite {
+	(options: TransformStyleWithViteOptions): Promise<{
+		code: string;
+		map: vite.TransformResult['map'];
+		deps: Set<string>;
+	} | null>;
 }
 
-/** Transform style using Vite hook */
-export async function transformWithVite({
-	value,
-	lang,
-	transformHook,
-	id,
-	ssr,
-	pluginContext,
-}: TransformWithViteOptions): Promise<vite.TransformResult | null> {
-	if (!STYLE_EXTENSIONS.has(lang)) {
-		return null; // only preprocess langs supported by Vite
-	}
-	return transformHook.call(pluginContext, value, id + `?astro&type=style&lang${lang}`, ssr);
+export function createTransformStyleWithViteFn(
+	viteConfig: vite.ResolvedConfig
+): TransformStyleWithVite {
+	const viteCSSPlugin = viteConfig.plugins.find(({ name }) => name === 'vite:css');
+	if (!viteCSSPlugin) throw new Error(`vite:css plugin couldn't be found`);
+	if (!viteCSSPlugin.transform) throw new Error(`vite:css has no transform() hook`);
+	const transformCss = viteCSSPlugin.transform as TransformHook;
+
+	return async function (
+		this: PluginContext,
+		{ id, source, lang, ssr, viteDevServer }: TransformStyleWithViteOptions
+	) {
+		if (!STYLE_EXTENSIONS.has(lang)) {
+			return null; // only preprocess langs supported by Vite
+		}
+
+		// Id must end with valid CSS extension for vite:css to process
+		const styleId = `${id}?astro&type=style&lang${lang}`;
+
+		viteDevServer?.moduleGraph.ensureEntryFromUrl(styleId, ssr, false);
+
+		const transformResult = await transformCss.call(this, source, styleId, ssr);
+
+		// NOTE: only `code` and `map` are returned by vite:css
+		const { code, map } = transformResult;
+		const deps = new Set<string>();
+
+		// Get deps from module created while transforming the styleId by Vite.
+		// In build, it's fine that we skip this as it's used by HMR only.
+		const mod = viteDevServer?.moduleGraph.getModuleById(styleId);
+		if (mod) {
+			// Get all @import references
+			for (const imported of mod.importedModules) {
+				if (imported.file) {
+					deps.add(imported.file);
+				}
+			}
+		}
+
+		return { code, map, deps };
+	};
 }
diff --git a/packages/astro/src/vite-plugin-markdown-legacy/index.ts b/packages/astro/src/vite-plugin-markdown-legacy/index.ts
index 5cb66caa5..1bfb8546b 100644
--- a/packages/astro/src/vite-plugin-markdown-legacy/index.ts
+++ b/packages/astro/src/vite-plugin-markdown-legacy/index.ts
@@ -4,13 +4,16 @@ import esbuild from 'esbuild';
 import fs from 'fs';
 import matter from 'gray-matter';
 import { fileURLToPath } from 'url';
-import type { Plugin } from 'vite';
+import type { Plugin, ViteDevServer } from 'vite';
 import type { AstroConfig } from '../@types/astro';
 import { pagesVirtualModuleId } from '../core/app/index.js';
 import { collectErrorMetadata } from '../core/errors.js';
 import type { LogOptions } from '../core/logger/core.js';
 import { cachedCompilation, CompileProps } from '../vite-plugin-astro/compile.js';
-import { getViteTransform, TransformHook } from '../vite-plugin-astro/styles.js';
+import {
+	createTransformStyleWithViteFn,
+	TransformStyleWithVite,
+} from '../vite-plugin-astro/styles.js';
 import type { PluginMetadata as AstroPluginMetadata } from '../vite-plugin-astro/types';
 import { getFileInfo } from '../vite-plugin-utils/index.js';
 
@@ -61,13 +64,14 @@ export default function markdown({ config, logging }: AstroPluginOptions): Plugi
 		return false;
 	}
 
-	let viteTransform: TransformHook;
+	let transformStyleWithVite: TransformStyleWithVite;
+	let viteDevServer: ViteDevServer | undefined;
 
 	return {
 		name: 'astro:markdown',
 		enforce: 'pre',
 		configResolved(_resolvedConfig) {
-			viteTransform = getViteTransform(_resolvedConfig);
+			transformStyleWithVite = createTransformStyleWithViteFn(_resolvedConfig);
 		},
 		async resolveId(id, importer, options) {
 			// Resolve any .md files with the `?content` cache buster. This should only come from
@@ -205,7 +209,8 @@ ${setup}`.trim();
 					moduleId: id,
 					source: astroResult,
 					ssr: Boolean(opts?.ssr),
-					viteTransform,
+					transformStyleWithVite,
+					viteDevServer,
 					pluginContext: this,
 				};