Add generic plugin for page-ssr injection (#4049)

* feat: add generic page-ssr plugin

* refactor: remove page-specific logic from astro/markdown/mdx plugins

* refactor: revert changes to vite-plugin-scripts

* fix: handle injected `page` scripts in build

* fix: prepend injected `page` scripts with `/@id/` in dev

Co-authored-by: Nate Moore <nate@astro.build>

.changeset/plenty-hats-grow.md

@@ -0,0 +1,6 @@
+---
+'astro': patch
+'@astrojs/mdx': patch
+---
+
+Improve `injectScript` handling for non-Astro pages

packages/astro/src/core/build/generate.ts

@@ -18,7 +18,7 @@ import {
 	removeTrailingForwardSlash,
 } from '../../core/path.js';
 import type { RenderOptions } from '../../core/render/core';
-import { BEFORE_HYDRATION_SCRIPT_ID } from '../../vite-plugin-scripts/index.js';
+import { BEFORE_HYDRATION_SCRIPT_ID, PAGE_SCRIPT_ID } from '../../vite-plugin-scripts/index.js';
 import { call as callEndpoint } from '../endpoint/index.js';
 import { debug, info } from '../logger/core.js';
 import { render } from '../render/core.js';
@@ -272,6 +272,18 @@ async function generatePath(
 	const links = createLinkStylesheetElementSet(linkIds.reverse(), site);
 	const scripts = createModuleScriptsSet(hoistedScripts ? [hoistedScripts] : [], site);
 
+	if (astroConfig._ctx.scripts.some((script) => script.stage === 'page')) {
+		const hashedFilePath = internals.entrySpecifierToBundleMap.get(PAGE_SCRIPT_ID);
+		if (typeof hashedFilePath !== 'string') {
+			throw new Error(`Cannot find the built path for ${PAGE_SCRIPT_ID}`);
+		}
+		const src = prependForwardSlash(npath.posix.join(astroConfig.base, hashedFilePath));
+		scripts.add({
+			props: { type: 'module', src },
+			children: '',
+		})
+	}
+
 	// Add all injected scripts to the page.
 	for (const script of astroConfig._ctx.scripts) {
 		if (script.stage === 'head-inline') {

packages/astro/src/core/build/static-build.ts

@@ -8,6 +8,7 @@ import { prependForwardSlash } from '../../core/path.js';
 import { emptyDir, isModeServerWithNoAdapter, removeDir } from '../../core/util.js';
 import { runHookBuildSetup } from '../../integrations/index.js';
 import { rollupPluginAstroBuildCSS } from '../../vite-plugin-build-css/index.js';
+import { PAGE_SCRIPT_ID } from '../../vite-plugin-scripts/index.js';
 import type { ViteConfigWithSSR } from '../create-vite';
 import { info } from '../logger/core.js';
 import { generatePages } from './generate.js';
@@ -85,6 +86,10 @@ Example:
 		...internals.discoveredScripts,
 	]);
 
+	if (astroConfig._ctx.scripts.some((script) => script.stage === 'page')) {
+		clientInput.add(PAGE_SCRIPT_ID);
+	}
+
 	// Run client build first, so the assets can be fed into the SSR rendered version.
 	timer.clientBuild = performance.now();
 	await clientBuild(opts, internals, clientInput);

packages/astro/src/core/build/vite-plugin-ssr.ts

@@ -8,7 +8,7 @@ import glob from 'fast-glob';
 import * as fs from 'fs';
 import { fileURLToPath } from 'url';
 import { runHookBuildSsr } from '../../integrations/index.js';
-import { BEFORE_HYDRATION_SCRIPT_ID } from '../../vite-plugin-scripts/index.js';
+import { BEFORE_HYDRATION_SCRIPT_ID, PAGE_SCRIPT_ID } from '../../vite-plugin-scripts/index.js';
 import { pagesVirtualModuleId } from '../app/index.js';
 import { serializeRouteData } from '../routing/index.js';
 import { addRollupInput } from './add-rollup-input.js';
@@ -123,12 +123,19 @@ function buildManifest(
 	const { astroConfig } = opts;
 
 	const routes: SerializedRouteInfo[] = [];
+	const entryModules = Object.fromEntries(internals.entrySpecifierToBundleMap.entries());
+	if (astroConfig._ctx.scripts.some((script) => script.stage === 'page')) {
+		staticFiles.push(entryModules[PAGE_SCRIPT_ID]);
+	}
 
 	for (const pageData of eachPageData(internals)) {
 		const scripts: SerializedRouteInfo['scripts'] = [];
 		if (pageData.hoistedScript) {
 			scripts.unshift(pageData.hoistedScript);
 		}
+		if (astroConfig._ctx.scripts.some((script) => script.stage === 'page')) {
+			scripts.push({ type: 'external', value: entryModules[PAGE_SCRIPT_ID] });
+		}
 
 		routes.push({
 			file: '',
@@ -144,7 +151,6 @@ function buildManifest(
 	}
 
 	// HACK! Patch this special one.
-	const entryModules = Object.fromEntries(internals.entrySpecifierToBundleMap.entries());
 	if (!(BEFORE_HYDRATION_SCRIPT_ID in entryModules)) {
 		entryModules[BEFORE_HYDRATION_SCRIPT_ID] =
 			'data:text/javascript;charset=utf-8,//[no before-hydration script]';

packages/astro/src/core/create-vite.ts

@@ -14,6 +14,7 @@ import astroIntegrationsContainerPlugin from '../vite-plugin-integrations-contai
 import jsxVitePlugin from '../vite-plugin-jsx/index.js';
 import markdownVitePlugin from '../vite-plugin-markdown/index.js';
 import astroScriptsPlugin from '../vite-plugin-scripts/index.js';
+import astroScriptsPageSSRPlugin from '../vite-plugin-scripts/page-ssr.js';
 import { createCustomViteLogger } from './errors.js';
 import { resolveDependency } from './util.js';
 
@@ -80,6 +81,7 @@ export async function createVite(
 			jsxVitePlugin({ config: astroConfig, logging }),
 			astroPostprocessVitePlugin({ config: astroConfig }),
 			astroIntegrationsContainerPlugin({ config: astroConfig }),
+			astroScriptsPageSSRPlugin({ config: astroConfig }),
 		],
 		publicDir: fileURLToPath(astroConfig.publicDir),
 		root: fileURLToPath(astroConfig.root),

packages/astro/src/core/render/dev/index.ts

@@ -9,6 +9,7 @@ import type {
 	SSRElement,
 	SSRLoadedRenderer,
 } from '../../../@types/astro';
+import { PAGE_SCRIPT_ID } from '../../../vite-plugin-scripts/index.js';
 import { prependForwardSlash } from '../../../core/path.js';
 import { LogOptions } from '../../logger/core.js';
 import { isPage } from '../../util.js';
@@ -124,13 +125,19 @@ export async function render(
 			children: '',
 		});
 	}
+
 	// TODO: We should allow adding generic HTML elements to the head, not just scripts
 	for (const script of astroConfig._ctx.scripts) {
 		if (script.stage === 'head-inline') {
 			scripts.add({
 				props: {},
 				children: script.content,
 			});
+		} else if (script.stage === 'page' && isPage(filePath, astroConfig)) {
+			scripts.add({
+				props: { type: 'module', src: `/@id/${PAGE_SCRIPT_ID}` },
+				children: '',
+			});
 		}
 	}
 

packages/astro/src/vite-plugin-astro/index.ts

@@ -9,8 +9,6 @@ import esbuild from 'esbuild';
 import slash from 'slash';
 import { fileURLToPath } from 'url';
 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, getCachedSource } from './compile.js';
 import { handleHotUpdate, trackCSSDependencies } from './hmr.js';
@@ -215,14 +213,6 @@ export default function astro({ config, logging }: AstroPluginOptions): vite.Plu
 			}
 
 			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,
@@ -269,10 +259,6 @@ export default function astro({ config, logging }: AstroPluginOptions): vite.Plu
 						i++;
 					}
 				}
-				// Add handling to inject scripts into each page JS bundle, if needed.
-				if (isPage) {
-					SUFFIX += `\nimport "${PAGE_SSR_SCRIPT_ID}";`;
-				}
 
 				// Prefer live reload to HMR in `.astro` files
 				if (!resolvedConfig.isProduction) {

packages/astro/src/vite-plugin-markdown/index.ts

@@ -10,11 +10,9 @@ import { pagesVirtualModuleId } from '../core/app/index.js';
 import { collectErrorMetadata } from '../core/errors.js';
 import type { LogOptions } from '../core/logger/core.js';
 import { warn } from '../core/logger/core.js';
-import { resolvePages } from '../core/util.js';
 import { cachedCompilation, CompileProps } from '../vite-plugin-astro/compile.js';
 import { getViteTransform, TransformHook } from '../vite-plugin-astro/styles.js';
 import type { PluginMetadata as AstroPluginMetadata } from '../vite-plugin-astro/types';
-import { PAGE_SSR_SCRIPT_ID } from '../vite-plugin-scripts/index.js';
 import { getFileInfo } from '../vite-plugin-utils/index.js';
 
 interface AstroPluginOptions {
@@ -147,8 +145,6 @@ export default function markdown({ config, logging }: AstroPluginOptions): Plugi
 				const isAstroFlavoredMd = config.legacy.astroFlavoredMarkdown;
 
 				const fileUrl = new URL(`file://${filename}`);
-				const isPage = fileUrl.pathname.startsWith(resolvePages(config).pathname);
-				const hasInjectedScript = isPage && config._ctx.scripts.some((s) => s.stage === 'page-ssr');
 
 				// Extract special frontmatter keys
 				let { data: frontmatter, content: markdownContent } = safeMatter(source, filename);
@@ -187,7 +183,6 @@ export default function markdown({ config, logging }: AstroPluginOptions): Plugi
 import Slugger from 'github-slugger';
 ${layout ? `import Layout from '${layout}';` : ''}
 ${isAstroFlavoredMd && components ? `import * from '${components}';` : ''}
-${hasInjectedScript ? `import '${PAGE_SSR_SCRIPT_ID}';` : ''}
 ${isAstroFlavoredMd ? setup : ''}
 
 const slugger = new Slugger();
@@ -224,7 +219,8 @@ ${isAstroFlavoredMd ? setup : ''}`.trim();
 				if (/\bLayout\b/.test(imports)) {
 					astroResult = `${prelude}\n<Layout content={$$content}>\n\n${astroResult}\n\n</Layout>`;
 				} else {
-					astroResult = `${prelude}\n${astroResult}`;
+					// Note: without a Layout, we need to inject `head` manually so `maybeRenderHead` runs
+					astroResult = `${prelude}\n<head></head>${astroResult}`;
 				}
 
 				// Transform from `.astro` to valid `.ts`

packages/astro/src/vite-plugin-scripts/page-ssr.ts

@@ -0,0 +1,50 @@
+import { Plugin as VitePlugin } from 'vite';
+import { AstroConfig } from '../@types/astro.js';
+import { PAGE_SSR_SCRIPT_ID } from './index.js';
+
+import { isPage } from '../core/util.js';
+import ancestor from 'common-ancestor-path';
+import MagicString from 'magic-string';
+
+export default function astroScriptsPostPlugin({ config }: { config: AstroConfig }): VitePlugin {
+	function normalizeFilename(filename: string) {
+		if (filename.startsWith('/@fs')) {
+			filename = filename.slice('/@fs'.length);
+		} else if (filename.startsWith('/') && !ancestor(filename, config.root.pathname)) {
+			filename = new URL('.' + filename, config.root).pathname;
+		}
+		return filename;
+	}
+
+	return {
+		name: 'astro:scripts:page-ssr',
+		enforce: 'post',
+
+		transform(this, code, id, options) {
+			if (!options?.ssr) return;
+
+			const hasInjectedScript = config._ctx.scripts.some((s) => s.stage === 'page-ssr');
+			if (!hasInjectedScript) return;
+
+			const filename = normalizeFilename(id);
+			let fileURL: URL;
+			try {
+				fileURL = new URL(`file://${filename}`);
+			} catch (e) {
+				// If we can't construct a valid URL, exit early
+				return;
+			}
+
+			const fileIsPage = isPage(fileURL, config);
+			if (!fileIsPage) return;
+
+			const s = new MagicString(code, { filename });
+			s.prepend(`import '${PAGE_SSR_SCRIPT_ID}';\n`);
+
+			return {
+				code: s.toString(),
+				map: s.generateMap(),
+			}
+		},
+	};
+}

packages/integrations/mdx/src/index.ts

@@ -125,13 +125,6 @@ export default function mdx(mdxOptions: MdxOptions = {}): AstroIntegration {
 									if (!id.endsWith('.mdx')) return;
 									const [, moduleExports] = parseESM(code);
 
-									// This adds support for injected "page-ssr" scripts in MDX files.
-									// TODO: This should only be happening on page entrypoints, not all imported MDX.
-									// TODO: This code is copy-pasted across all Astro/Vite plugins that deal with page
-									// entrypoints (.astro, .md, .mdx). This should be handled in some centralized place,
-									// or otherwise refactored to not require copy-paste handling logic.
-									code += `\nimport "${'astro:scripts/page-ssr.js'}";`;
-
 									const { fileUrl, fileId } = getFileInfo(id, config);
 									if (!moduleExports.includes('url')) {
 										code += `\nexport const url = ${JSON.stringify(fileUrl)};`;