Allow defining Astro components in Vite plugins (#3889)

* Allow defining Astro components in Vite plugins

* Adds a changeset

* Move non-main compilation into load

* Use the cachedCompilation in the markdown plugin

* Fix HMR test

* Simplify getNormalizedID

* Use a windows-friendly virtual module id for the test

11 files changed, 194 insertions, 40 deletions

diff --git a/.changeset/hungry-cougars-yell.md b/.changeset/hungry-cougars-yell.md
new file mode 100644
index 000000000..658e0bb91
--- /dev/null
+++ b/.changeset/hungry-cougars-yell.md
@@ -0,0 +1,5 @@
+---
+'astro': patch
+---
+
+Allow defining Astro components in Vite plugins
diff --git a/packages/astro/src/vite-plugin-astro/compile.ts b/packages/astro/src/vite-plugin-astro/compile.ts
index 7b8c6c7fd..5a1d71868 100644
--- a/packages/astro/src/vite-plugin-astro/compile.ts
+++ b/packages/astro/src/vite-plugin-astro/compile.ts
@@ -11,7 +11,10 @@ import { viteID } from '../core/util.js';
 import { transformWithVite } from './styles.js';
 
 type CompilationCache = Map<string, CompileResult>;
-type CompileResult = TransformResult & { rawCSSDeps: Set<string> };
+type CompileResult = TransformResult & {
+	rawCSSDeps: Set<string>;
+	source: string;
+};
 
 /**
  * Note: this is currently needed because Astro is directly using a Vite internal CSS transform. This gives us
@@ -44,6 +47,16 @@ export interface CompileProps {
 	pluginContext: PluginContext;
 }
 
+function getNormalizedID(filename: string): string {
+	try {
+		const filenameURL = new URL(`file://${filename}`);
+		return fileURLToPath(filenameURL);
+	} catch(err) {
+		// Not a real file, so just use the provided filename as the normalized id
+		return filename;
+	}
+}
+
 async function compile({
 	config,
 	filename,
@@ -53,9 +66,7 @@ async function compile({
 	viteTransform,
 	pluginContext,
 }: CompileProps): Promise<CompileResult> {
-	const filenameURL = new URL(`file://${filename}`);
-	const normalizedID = fileURLToPath(filenameURL);
-
+	const normalizedID = getNormalizedID(filename);
 	let rawCSSDeps = new Set<string>();
 	let cssTransformError: Error | undefined;
 
@@ -141,6 +152,9 @@ async function compile({
 		rawCSSDeps: {
 			value: rawCSSDeps,
 		},
+		source: {
+			value: source,
+		},
 	});
 
 	return compileResult;
@@ -150,6 +164,13 @@ export function isCached(config: AstroConfig, filename: string) {
 	return configCache.has(config) && configCache.get(config)!.has(filename);
 }
 
+export function getCachedSource(config: AstroConfig, filename: string): string | null {
+	if(!isCached(config, filename)) return null;
+	let src = configCache.get(config)!.get(filename);
+	if(!src) return null;
+	return src.source;
+}
+
 export function invalidateCompilation(config: AstroConfig, filename: string) {
 	if (configCache.has(config)) {
 		const cache = configCache.get(config)!;
diff --git a/packages/astro/src/vite-plugin-astro/index.ts b/packages/astro/src/vite-plugin-astro/index.ts
index 1318d7458..195bb149d 100644
--- a/packages/astro/src/vite-plugin-astro/index.ts
+++ b/packages/astro/src/vite-plugin-astro/index.ts
@@ -13,7 +13,7 @@ import { isRelativePath, startsWithForwardSlash } from '../core/path.js';
 import { resolvePages } from '../core/util.js';
 import { PAGE_SCRIPT_ID, PAGE_SSR_SCRIPT_ID } from '../vite-plugin-scripts/index.js';
 import { getFileInfo } from '../vite-plugin-utils/index.js';
-import { cachedCompilation, CompileProps } from './compile.js';
+import { cachedCompilation, CompileProps, getCachedSource } from './compile.js';
 import { handleHotUpdate, trackCSSDependencies } from './hmr.js';
 import { parseAstroRequest, ParsedRequestResult } from './query.js';
 import { getViteTransform, TransformHook } from './styles.js';
@@ -96,24 +96,25 @@ export default function astro({ config, logging }: AstroPluginOptions): vite.Plu
 				return id;
 			}
 		},
-		async load(this: PluginContext, id, opts) {
+		async load(id, opts) {
 			const parsedId = parseAstroRequest(id);
 			const query = parsedId.query;
-			if (!id.endsWith('.astro') && !query.astro) {
+			if (!query.astro) {
 				return null;
 			}
-			// if we still get a relative path here, vite couldn't resolve the import
-			if (isRelativePath(parsedId.filename)) {
+			let filename = parsedId.filename;
+			// For CSS / hoisted scripts we need to load the source ourselves.
+			// It should be in the compilation cache at this point.
+			let raw = await this.resolve(filename, undefined);
+			if(!raw) {
 				return null;
 			}
 
-			const filename = normalizeFilename(parsedId.filename);
-			const fileUrl = new URL(`file://${filename}`);
-			let source = await fs.promises.readFile(fileUrl, 'utf-8');
-			const isPage = fileUrl.pathname.startsWith(resolvePages(config).pathname);
-			if (isPage && config._ctx.scripts.some((s) => s.stage === 'page')) {
-				source += `\n<script src="${PAGE_SCRIPT_ID}" />`;
+			let source = getCachedSource(config, raw.id);
+			if(!source) {
+				return null;
 			}
+
 			const compileProps: CompileProps = {
 				config,
 				filename,
@@ -123,14 +124,15 @@ export default function astro({ config, logging }: AstroPluginOptions): vite.Plu
 				viteTransform,
 				pluginContext: this,
 			};
-			if (query.astro) {
-				if (query.type === 'style') {
+
+			switch(query.type) {
+				case 'style': {
 					if (typeof query.index === 'undefined') {
 						throw new Error(`Requests for Astro CSS must include an index.`);
 					}
-
+	
 					const transformResult = await cachedCompilation(compileProps);
-
+	
 					// Track any CSS dependencies so that HMR is triggered when they change.
 					await trackCSSDependencies.call(this, {
 						viteDevServer,
@@ -140,11 +142,12 @@ export default function astro({ config, logging }: AstroPluginOptions): vite.Plu
 					});
 					const csses = transformResult.css;
 					const code = csses[query.index];
-
+	
 					return {
 						code,
 					};
-				} else if (query.type === 'script') {
+				}
+				case 'script': {
 					if (typeof query.index === 'undefined') {
 						throw new Error(`Requests for hoisted scripts must include an index`);
 					}
@@ -154,15 +157,15 @@ export default function astro({ config, logging }: AstroPluginOptions): vite.Plu
 							code: `/* client hoisted script, empty in SSR: ${id} */`,
 						};
 					}
-
+	
 					const transformResult = await cachedCompilation(compileProps);
 					const scripts = transformResult.scripts;
 					const hoistedScript = scripts[query.index];
-
+	
 					if (!hoistedScript) {
 						throw new Error(`No hoisted script at index ${query.index}`);
 					}
-
+	
 					if (hoistedScript.type === 'external') {
 						const src = hoistedScript.src!;
 						if (src.startsWith('/') && !isBrowserPath(src)) {
@@ -172,7 +175,7 @@ export default function astro({ config, logging }: AstroPluginOptions): vite.Plu
 							);
 						}
 					}
-
+	
 					return {
 						code:
 							hoistedScript.type === 'inline'
@@ -185,7 +188,39 @@ export default function astro({ config, logging }: AstroPluginOptions): vite.Plu
 						},
 					};
 				}
+				default: return null;
 			}
+		},
+		async transform(this: PluginContext, source, id, opts) {
+			const parsedId = parseAstroRequest(id);
+			const query = parsedId.query;
+			if (!id.endsWith('.astro') || query.astro) {
+				return source;
+			}
+			// if we still get a relative path here, vite couldn't resolve the import
+			if (isRelativePath(parsedId.filename)) {
+				return source;
+			}
+
+			const filename = normalizeFilename(parsedId.filename);
+			let isPage = false;
+			try {
+				const fileUrl = new URL(`file://${filename}`);
+				isPage = fileUrl.pathname.startsWith(resolvePages(config).pathname);
+			} catch {}
+			if (isPage && config._ctx.scripts.some((s) => s.stage === 'page')) {
+				source += `\n<script src="${PAGE_SCRIPT_ID}" />`;
+			}
+			const compileProps: CompileProps = {
+				config,
+				filename,
+				moduleId: id,
+				source,
+				ssr: Boolean(opts?.ssr),
+				viteTransform,
+				pluginContext: this,
+			};
+			
 
 			try {
 				const transformResult = await cachedCompilation(compileProps);
diff --git a/packages/astro/src/vite-plugin-markdown/index.ts b/packages/astro/src/vite-plugin-markdown/index.ts
index 8e0410bd7..e76afbf7b 100644
--- a/packages/astro/src/vite-plugin-markdown/index.ts
+++ b/packages/astro/src/vite-plugin-markdown/index.ts
@@ -12,6 +12,8 @@ import { collectErrorMetadata } from '../core/errors.js';
 import { prependForwardSlash } from '../core/path.js';
 import { resolvePages, viteID } from '../core/util.js';
 import type { PluginMetadata as AstroPluginMetadata } from '../vite-plugin-astro/types';
+import { cachedCompilation, CompileProps } from '../vite-plugin-astro/compile.js';
+import { getViteTransform, TransformHook } from '../vite-plugin-astro/styles.js';
 import { PAGE_SSR_SCRIPT_ID } from '../vite-plugin-scripts/index.js';
 import { getFileInfo } from '../vite-plugin-utils/index.js';
 
@@ -61,9 +63,14 @@ export default function markdown({ config }: AstroPluginOptions): Plugin {
 		return false;
 	}
 
+	let viteTransform: TransformHook;
+
 	return {
 		name: 'astro:markdown',
 		enforce: 'pre',
+		configResolved(_resolvedConfig) {
+			viteTransform = getViteTransform(_resolvedConfig);
+		},
 		async resolveId(id, importer, options) {
 			// Resolve any .md files with the `?content` cache buster. This should only come from
 			// an already-resolved JS module wrapper. Needed to prevent infinite loops in Vite.
@@ -85,7 +92,7 @@ export default function markdown({ config }: AstroPluginOptions): Plugin {
 			// In all other cases, we do nothing and rely on normal Vite resolution.
 			return undefined;
 		},
-		async load(id) {
+		async load(id, opts) {
 			// A markdown file has been imported via ESM!
 			// Return the file's JS representation, including all Markdown
 			// frontmatter and a deferred `import() of the compiled markdown content.
@@ -174,21 +181,17 @@ ${setup}`.trim();
 				}
 
 				// Transform from `.astro` to valid `.ts`
-				let transformResult = await transform(astroResult, {
-					pathname: '/@fs' + prependForwardSlash(fileUrl.pathname),
-					projectRoot: config.root.toString(),
-					site: config.site
-						? new URL(config.base, config.site).toString()
-						: `http://localhost:${config.server.port}/`,
-					sourcefile: id,
-					sourcemap: 'inline',
-					// TODO: baseline flag
-					experimentalStaticExtraction: true,
-					internalURL: `/@fs${prependForwardSlash(
-						viteID(new URL('../runtime/server/index.js', import.meta.url))
-					)}`,
-				});
+				const compileProps: CompileProps = {
+					config,
+					filename,
+					moduleId: id,
+					source: astroResult,
+					ssr: Boolean(opts?.ssr),
+					viteTransform,
+					pluginContext: this,
+				};
 
+				let transformResult = await cachedCompilation(compileProps)
 				let { code: tsResult } = transformResult;
 
 				tsResult = `\nexport const metadata = ${JSON.stringify(metadata)};
diff --git a/packages/astro/test/fixtures/virtual-astro-file/astro.config.mjs b/packages/astro/test/fixtures/virtual-astro-file/astro.config.mjs
new file mode 100644
index 000000000..2f864ac37
--- /dev/null
+++ b/packages/astro/test/fixtures/virtual-astro-file/astro.config.mjs
@@ -0,0 +1,9 @@
+import { defineConfig } from 'astro/config';
+import myPlugin from './src/plugin/my-plugin.mjs';
+
+// https://astro.build/config
+export default defineConfig({
+	vite: {
+		plugins: [myPlugin()]
+	}
+});
diff --git a/packages/astro/test/fixtures/virtual-astro-file/package.json b/packages/astro/test/fixtures/virtual-astro-file/package.json
new file mode 100644
index 000000000..40f712e54
--- /dev/null
+++ b/packages/astro/test/fixtures/virtual-astro-file/package.json
@@ -0,0 +1,8 @@
+{
+  "name": "@test/virtual-astro-file",
+  "version": "0.0.0",
+  "private": true,
+  "dependencies": {
+    "astro": "workspace:*"
+  }
+}
diff --git a/packages/astro/test/fixtures/virtual-astro-file/src/pages/index.astro b/packages/astro/test/fixtures/virtual-astro-file/src/pages/index.astro
new file mode 100644
index 000000000..e480002a9
--- /dev/null
+++ b/packages/astro/test/fixtures/virtual-astro-file/src/pages/index.astro
@@ -0,0 +1,10 @@
+---
+import Something from '@my-plugin/virtual.astro';
+---
+<html>
+	<head><title>Testing</title></head>
+	<body>
+		<Something />
+	</body>
+</html>
+
diff --git a/packages/astro/test/fixtures/virtual-astro-file/src/plugin/my-plugin.mjs b/packages/astro/test/fixtures/virtual-astro-file/src/plugin/my-plugin.mjs
new file mode 100644
index 000000000..e01f6a0fd
--- /dev/null
+++ b/packages/astro/test/fixtures/virtual-astro-file/src/plugin/my-plugin.mjs
@@ -0,0 +1,27 @@
+
+
+export default function myPlugin() {
+	const pluginId = `@my-plugin/virtual.astro`;
+  return {
+    enforce: 'pre',
+    name: 'virtual-astro-plugin',
+    resolveId(id) {
+      if (id === pluginId) return id;
+    },
+    load(id) {
+      if (id === pluginId) {
+        return `---
+const works = true;
+---
+<h1 id="something">This is a virtual module id</h1>
+<h2 id="works">{works}</h2>
+<style>
+  h1 {
+		color: green;
+	}
+</style>
+`;
+      }
+    },
+  };
+}
diff --git a/packages/astro/test/hmr-css.test.js b/packages/astro/test/hmr-css.test.js
index f8f741904..b2b4341b4 100644
--- a/packages/astro/test/hmr-css.test.js
+++ b/packages/astro/test/hmr-css.test.js
@@ -22,6 +22,9 @@ describe('HMR - CSS', () => {
 	});
 
 	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'
 		);
diff --git a/packages/astro/test/virtual-astro-file.js b/packages/astro/test/virtual-astro-file.js
new file mode 100644
index 000000000..e5247a684
--- /dev/null
+++ b/packages/astro/test/virtual-astro-file.js
@@ -0,0 +1,27 @@
+import { expect } from 'chai';
+import * as cheerio from 'cheerio';
+import { loadFixture } from './test-utils.js';
+
+describe('Loading virtual Astro files', () => {
+	let fixture;
+
+	before(async () => {
+		fixture = await loadFixture({ root: './fixtures/virtual-astro-file/' });
+		await fixture.build();
+	});
+
+	it('renders the component', async () => {
+		const html = await fixture.readFile('/index.html');
+		const $ = cheerio.load(html);
+		expect($('#something')).to.have.a.lengthOf(1);
+		expect($('#works').text()).to.equal('true');
+	});
+
+	it('builds component CSS', async () => {
+		const html = await fixture.readFile('/index.html');
+		const $ = cheerio.load(html);
+		const href = $('link').attr('href');
+		const css = await fixture.readFile(href);
+		expect(css).to.match(/green/, 'css bundled from virtual astro module');
+	});
+});
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
index 16cc5280d..2a61a77b8 100644
--- a/pnpm-lock.yaml
+++ b/pnpm-lock.yaml
@@ -1802,6 +1802,12 @@ importers:
       postcss: 8.4.14
       tailwindcss: 3.1.5
 
+  packages/astro/test/fixtures/virtual-astro-file:
+    specifiers:
+      astro: workspace:*
+    dependencies:
+      astro: link:../../..
+
   packages/astro/test/fixtures/vue-component:
     specifiers:
       '@astrojs/vue': workspace:*