From 8e3e4894c9c7125bad74496761dfde8551ec9a2a Mon Sep 17 00:00:00 2001
From: Nate Moore <natemoo-re@users.noreply.github.com>
Date: Mon, 20 Jun 2022 14:05:16 -0500
Subject: [PATCH] Add private `addPageExtension` hook (#3628)

* feat: add private `addPageExtensions` hook

* chore: remove renderer binding

Co-authored-by: Nate Moore <nate@astro.build>
---
 packages/astro/src/@types/astro.ts            |  1 +
 packages/astro/src/core/build/static-build.ts | 70 +++++++++++--------
 packages/astro/src/core/config.ts             |  2 +-
 packages/astro/src/core/render/core.ts        | 16 +++--
 .../astro/src/core/routing/manifest/create.ts |  2 +-
 packages/astro/src/integrations/index.ts      | 18 ++++-
 .../astro/src/runtime/server/hydration.ts     |  7 ++
 packages/astro/src/runtime/server/index.ts    |  5 +-
 .../astro.config.mjs                          |  6 ++
 .../integration.js                            | 10 +++
 .../package.json                              |  9 +++
 .../src/components/test.astro                 |  1 +
 .../src/pages/test.mjs                        |  3 +
 .../integration-add-page-extension.test.js    | 19 +++++
 pnpm-lock.yaml                                |  6 ++
 15 files changed, 134 insertions(+), 41 deletions(-)
 create mode 100644 packages/astro/test/fixtures/integration-add-page-extension/astro.config.mjs
 create mode 100644 packages/astro/test/fixtures/integration-add-page-extension/integration.js
 create mode 100644 packages/astro/test/fixtures/integration-add-page-extension/package.json
 create mode 100644 packages/astro/test/fixtures/integration-add-page-extension/src/components/test.astro
 create mode 100644 packages/astro/test/fixtures/integration-add-page-extension/src/pages/test.mjs
 create mode 100644 packages/astro/test/integration-add-page-extension.test.js

diff --git a/packages/astro/src/@types/astro.ts b/packages/astro/src/@types/astro.ts
index 5c398566c..dc2af7db3 100644
--- a/packages/astro/src/@types/astro.ts
+++ b/packages/astro/src/@types/astro.ts
@@ -726,6 +726,7 @@ export interface AstroConfig extends z.output<typeof AstroConfigSchema> {
 	// that is different from the user-exposed configuration.
 	// TODO: Create an AstroConfig class to manage this, long-term.
 	_ctx: {
+		pageExtensions: string[];
 		injectedRoutes: InjectedRoute[];
 		adapter: AstroAdapter | undefined;
 		renderers: AstroRenderer[];
diff --git a/packages/astro/src/core/build/static-build.ts b/packages/astro/src/core/build/static-build.ts
index c242c98b3..2fe43e6c8 100644
--- a/packages/astro/src/core/build/static-build.ts
+++ b/packages/astro/src/core/build/static-build.ts
@@ -57,46 +57,54 @@ export async function staticBuild(opts: StaticBuildOptions) {
 			const [renderers, mod] = pageData.preload;
 			const metadata = mod.$$metadata;
 
-			// Track client:only usage so we can map their CSS back to the Page they are used in.
-			const clientOnlys = Array.from(metadata.clientOnlyComponentPaths());
-			trackClientOnlyPageDatas(internals, pageData, clientOnlys);
-
 			const topLevelImports = new Set([
-				// Any component that gets hydrated
-				// 'components/Counter.jsx'
-				// { 'components/Counter.jsx': 'counter.hash.js' }
-				...metadata.hydratedComponentPaths(),
-				// Client-only components
-				...clientOnlys,
 				// The client path for each renderer
 				...renderers
 					.filter((renderer) => !!renderer.clientEntrypoint)
 					.map((renderer) => renderer.clientEntrypoint!),
 			]);
 
-			// Add hoisted scripts
-			const hoistedScripts = new Set(metadata.hoistedScriptPaths());
-			if (hoistedScripts.size) {
-				const uniqueHoistedId = JSON.stringify(Array.from(hoistedScripts).sort());
-				let moduleId: string;
-
-				// If we're already tracking this set of hoisted scripts, get the unique id
-				if (uniqueHoistedIds.has(uniqueHoistedId)) {
-					moduleId = uniqueHoistedIds.get(uniqueHoistedId)!;
-				} else {
-					// Otherwise, create a unique id for this set of hoisted scripts
-					moduleId = `/astro/hoisted.js?q=${uniqueHoistedIds.size}`;
-					uniqueHoistedIds.set(uniqueHoistedId, moduleId);
+			if (metadata) {
+				// Any component that gets hydrated
+				// 'components/Counter.jsx'
+				// { 'components/Counter.jsx': 'counter.hash.js' }
+				for (const hydratedComponentPath of metadata.hydratedComponentPaths()) {
+					topLevelImports.add(hydratedComponentPath);
 				}
-				topLevelImports.add(moduleId);
 
-				// Make sure to track that this page uses this set of hoisted scripts
-				if (internals.hoistedScriptIdToPagesMap.has(moduleId)) {
-					const pages = internals.hoistedScriptIdToPagesMap.get(moduleId);
-					pages!.add(astroModuleId);
-				} else {
-					internals.hoistedScriptIdToPagesMap.set(moduleId, new Set([astroModuleId]));
-					internals.hoistedScriptIdToHoistedMap.set(moduleId, hoistedScripts);
+				// Track client:only usage so we can map their CSS back to the Page they are used in.
+				const clientOnlys = Array.from(metadata.clientOnlyComponentPaths());
+				trackClientOnlyPageDatas(internals, pageData, clientOnlys);
+				
+				// Client-only components
+				for (const clientOnly of clientOnlys) {
+					topLevelImports.add(clientOnly)
+				}
+
+				// Add hoisted scripts
+				const hoistedScripts = new Set(metadata.hoistedScriptPaths());
+				if (hoistedScripts.size) {
+					const uniqueHoistedId = JSON.stringify(Array.from(hoistedScripts).sort());
+					let moduleId: string;
+
+					// If we're already tracking this set of hoisted scripts, get the unique id
+					if (uniqueHoistedIds.has(uniqueHoistedId)) {
+						moduleId = uniqueHoistedIds.get(uniqueHoistedId)!;
+					} else {
+						// Otherwise, create a unique id for this set of hoisted scripts
+						moduleId = `/astro/hoisted.js?q=${uniqueHoistedIds.size}`;
+						uniqueHoistedIds.set(uniqueHoistedId, moduleId);
+					}
+					topLevelImports.add(moduleId);
+
+					// Make sure to track that this page uses this set of hoisted scripts
+					if (internals.hoistedScriptIdToPagesMap.has(moduleId)) {
+						const pages = internals.hoistedScriptIdToPagesMap.get(moduleId);
+						pages!.add(astroModuleId);
+					} else {
+						internals.hoistedScriptIdToPagesMap.set(moduleId, new Set([astroModuleId]));
+						internals.hoistedScriptIdToHoistedMap.set(moduleId, hoistedScripts);
+					}
 				}
 			}
 
diff --git a/packages/astro/src/core/config.ts b/packages/astro/src/core/config.ts
index 5ef15a1af..8388cbf0b 100644
--- a/packages/astro/src/core/config.ts
+++ b/packages/astro/src/core/config.ts
@@ -338,7 +338,7 @@ export async function validateConfig(
 	// First-Pass Validation
 	const result = {
 		...(await AstroConfigRelativeSchema.parseAsync(userConfig)),
-		_ctx: { scripts: [], renderers: [], injectedRoutes: [], adapter: undefined },
+		_ctx: { pageExtensions: [], scripts: [], renderers: [], injectedRoutes: [], adapter: undefined },
 	};
 	// Final-Pass Validation (perform checks that require the full config object)
 	if (
diff --git a/packages/astro/src/core/render/core.ts b/packages/astro/src/core/render/core.ts
index 61b46b403..6332fbc09 100644
--- a/packages/astro/src/core/render/core.ts
+++ b/packages/astro/src/core/render/core.ts
@@ -9,7 +9,7 @@ import type {
 } from '../../@types/astro';
 import type { LogOptions } from '../logger/core.js';
 
-import { renderHead, renderPage } from '../../runtime/server/index.js';
+import { renderHead, renderPage, renderComponent } from '../../runtime/server/index.js';
 import { getParams } from '../routing/params.js';
 import { createResult } from './result.js';
 import { callGetStaticPaths, findPathItemByKey, RouteCache } from './route-cache.js';
@@ -126,8 +126,6 @@ export async function render(
 	const Component = await mod.default;
 	if (!Component)
 		throw new Error(`Expected an exported Astro component but received typeof ${typeof Component}`);
-	if (!Component.isAstroComponentFactory)
-		throw new Error(`Unable to SSR non-Astro component (${route?.component})`);
 
 	const result = createResult({
 		links,
@@ -146,7 +144,17 @@ export async function render(
 		ssr,
 	});
 
-	let page = await renderPage(result, Component, pageProps, null);
+	let page: Awaited<ReturnType<typeof renderPage>>;
+	if (!Component.isAstroComponentFactory) {
+		const props: Record<string, any> = { ...(pageProps ?? {}), 'server:root': true };
+		const html = await renderComponent(result, Component.name, Component, props, null);
+		page = {
+			type: 'html',
+			html: html.toString()
+		}
+	} else {
+		page = await renderPage(result, Component, pageProps, null);
+	}
 
 	if (page.type === 'response') {
 		return page;
diff --git a/packages/astro/src/core/routing/manifest/create.ts b/packages/astro/src/core/routing/manifest/create.ts
index 85e5fea15..b2ce27062 100644
--- a/packages/astro/src/core/routing/manifest/create.ts
+++ b/packages/astro/src/core/routing/manifest/create.ts
@@ -165,7 +165,7 @@ export function createRouteManifest(
 ): ManifestData {
 	const components: string[] = [];
 	const routes: RouteData[] = [];
-	const validPageExtensions: Set<string> = new Set(['.astro', '.md']);
+	const validPageExtensions: Set<string> = new Set(['.astro', '.md', ...config._ctx.pageExtensions]);
 	const validEndpointExtensions: Set<string> = new Set(['.js', '.ts']);
 
 	function walk(dir: string, parentSegments: RoutePart[][], parentParams: string[]) {
diff --git a/packages/astro/src/integrations/index.ts b/packages/astro/src/integrations/index.ts
index a1ec05f6a..a519b0516 100644
--- a/packages/astro/src/integrations/index.ts
+++ b/packages/astro/src/integrations/index.ts
@@ -1,6 +1,6 @@
 import type { AddressInfo } from 'net';
 import type { ViteDevServer } from 'vite';
-import { AstroConfig, AstroRenderer, BuildConfig, RouteData } from '../@types/astro.js';
+import { AstroConfig, AstroIntegration, AstroRenderer, BuildConfig, RouteData } from '../@types/astro.js';
 import ssgAdapter from '../adapter-ssg/index.js';
 import type { SerializedSSRManifest } from '../core/app/types';
 import type { PageBuildData } from '../core/build/types';
@@ -8,6 +8,8 @@ import { mergeConfig } from '../core/config.js';
 import type { ViteConfigWithSSR } from '../core/create-vite.js';
 import { isBuildingToSSR } from '../core/util.js';
 
+type Hooks<Hook extends keyof AstroIntegration['hooks'], Fn = AstroIntegration['hooks'][Hook]> = Fn extends (...args: any) => any ? Parameters<Fn>[0] : never;
+
 export async function runHookConfigSetup({
 	config: _config,
 	command,
@@ -34,7 +36,7 @@ export async function runHookConfigSetup({
 		 * ```
 		 */
 		if (integration?.hooks?.['astro:config:setup']) {
-			await integration.hooks['astro:config:setup']({
+			const hooks: Hooks<'astro:config:setup'> = {
 				config: updatedConfig,
 				command,
 				addRenderer(renderer: AstroRenderer) {
@@ -49,7 +51,17 @@ export async function runHookConfigSetup({
 				injectRoute: (injectRoute) => {
 					updatedConfig._ctx.injectedRoutes.push(injectRoute);
 				},
-			});
+			}
+			// Semi-private `addPageExtension` hook
+			Object.defineProperty(hooks, 'addPageExtension', {
+				value: (...input: (string|string[])[]) => {
+					const exts = (input.flat(Infinity) as string[]).map(ext => `.${ext.replace(/^\./, '')}`);
+					updatedConfig._ctx.pageExtensions.push(...exts);
+				},
+				writable: false,
+				enumerable: false
+			})
+			await integration.hooks['astro:config:setup'](hooks);
 		}
 	}
 	return updatedConfig;
diff --git a/packages/astro/src/runtime/server/hydration.ts b/packages/astro/src/runtime/server/hydration.ts
index ef677c8b8..7f3181484 100644
--- a/packages/astro/src/runtime/server/hydration.ts
+++ b/packages/astro/src/runtime/server/hydration.ts
@@ -11,6 +11,7 @@ import { serializeListValue } from './util.js';
 const HydrationDirectives = ['load', 'idle', 'media', 'visible', 'only'];
 
 interface ExtractedProps {
+	isPage: boolean;
 	hydration: {
 		directive: string;
 		value: string;
@@ -24,10 +25,16 @@ interface ExtractedProps {
 // Finds these special props and removes them from what gets passed into the component.
 export function extractDirectives(inputProps: Record<string | number, any>): ExtractedProps {
 	let extracted: ExtractedProps = {
+		isPage: false,
 		hydration: null,
 		props: {},
 	};
 	for (const [key, value] of Object.entries(inputProps)) {
+		if (key.startsWith('server:')) {
+			if (key === 'server:root') {
+				extracted.isPage = true;
+			}
+		}
 		if (key.startsWith('client:')) {
 			if (!extracted.hydration) {
 				extracted.hydration = {
diff --git a/packages/astro/src/runtime/server/index.ts b/packages/astro/src/runtime/server/index.ts
index 539bfad63..fecf3317f 100644
--- a/packages/astro/src/runtime/server/index.ts
+++ b/packages/astro/src/runtime/server/index.ts
@@ -181,7 +181,7 @@ export async function renderComponent(
 	const { renderers } = result._metadata;
 	const metadata: AstroComponentMetadata = { displayName };
 
-	const { hydration, props } = extractDirectives(_props);
+	const { hydration, isPage, props } = extractDirectives(_props);
 	let html = '';
 	let needsHydrationScript = hydration && determineIfNeedsHydrationScript(result);
 	let needsDirectiveScript =
@@ -317,6 +317,9 @@ If you're still stuck, please open an issue on GitHub or join us at https://astr
 	}
 
 	if (!hydration) {
+		if (isPage) {
+			return html;
+		}
 		return markHTMLString(html.replace(/\<\/?astro-fragment\>/g, ''));
 	}
 
diff --git a/packages/astro/test/fixtures/integration-add-page-extension/astro.config.mjs b/packages/astro/test/fixtures/integration-add-page-extension/astro.config.mjs
new file mode 100644
index 000000000..0a0a33697
--- /dev/null
+++ b/packages/astro/test/fixtures/integration-add-page-extension/astro.config.mjs
@@ -0,0 +1,6 @@
+import { defineConfig } from 'rollup'
+import test from './integration.js'
+
+export default defineConfig({
+	integrations: [test()]
+})
diff --git a/packages/astro/test/fixtures/integration-add-page-extension/integration.js b/packages/astro/test/fixtures/integration-add-page-extension/integration.js
new file mode 100644
index 000000000..8050a061d
--- /dev/null
+++ b/packages/astro/test/fixtures/integration-add-page-extension/integration.js
@@ -0,0 +1,10 @@
+export default function() {
+	return {
+		name: '@astrojs/test-integration',
+		hooks: {
+			'astro:config:setup': ({ addPageExtension }) => {
+				addPageExtension('.mjs')
+			}
+		}
+	}
+}
diff --git a/packages/astro/test/fixtures/integration-add-page-extension/package.json b/packages/astro/test/fixtures/integration-add-page-extension/package.json
new file mode 100644
index 000000000..cae9492df
--- /dev/null
+++ b/packages/astro/test/fixtures/integration-add-page-extension/package.json
@@ -0,0 +1,9 @@
+{
+  "name": "@test/integration-add-page-extension",
+  "type": "module",
+  "version": "0.0.0",
+  "private": true,
+  "dependencies": {
+    "astro": "workspace:*"
+  }
+}
diff --git a/packages/astro/test/fixtures/integration-add-page-extension/src/components/test.astro b/packages/astro/test/fixtures/integration-add-page-extension/src/components/test.astro
new file mode 100644
index 000000000..597ecf5fc
--- /dev/null
+++ b/packages/astro/test/fixtures/integration-add-page-extension/src/components/test.astro
@@ -0,0 +1 @@
+<h1>Hello world!</h1>
diff --git a/packages/astro/test/fixtures/integration-add-page-extension/src/pages/test.mjs b/packages/astro/test/fixtures/integration-add-page-extension/src/pages/test.mjs
new file mode 100644
index 000000000..b6bed7c56
--- /dev/null
+++ b/packages/astro/test/fixtures/integration-add-page-extension/src/pages/test.mjs
@@ -0,0 +1,3 @@
+// Convulted test case, rexport astro file from new `.mjs` page
+import Test from '../components/test.astro';
+export default Test;
diff --git a/packages/astro/test/integration-add-page-extension.test.js b/packages/astro/test/integration-add-page-extension.test.js
new file mode 100644
index 000000000..28c11d4fc
--- /dev/null
+++ b/packages/astro/test/integration-add-page-extension.test.js
@@ -0,0 +1,19 @@
+import { expect } from 'chai';
+import * as cheerio from 'cheerio';
+import { loadFixture } from './test-utils.js';
+
+describe('Integration addPageExtension', () => {
+	/** @type {import('./test-utils').Fixture} */
+	let fixture;
+
+	before(async () => {
+		fixture = await loadFixture({ root: './fixtures/integration-add-page-extension/' });
+		await fixture.build();
+	});
+
+	it('supports .mjs files', async () => {
+		const html = await fixture.readFile('/test/index.html');
+		const $ = cheerio.load(html);
+		expect($('h1').text()).to.equal('Hello world!');
+	});
+});
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
index 7e862156e..c2752da36 100644
--- a/pnpm-lock.yaml
+++ b/pnpm-lock.yaml
@@ -1388,6 +1388,12 @@ importers:
       '@fontsource/montserrat': 4.5.11
       astro: link:../../..
 
+  packages/astro/test/fixtures/integration-add-page-extension:
+    specifiers:
+      astro: workspace:*
+    dependencies:
+      astro: link:../../..
+
   packages/astro/test/fixtures/legacy-build:
     specifiers:
       '@astrojs/vue': workspace:*