Simplified head injection (#6034)

* Simplified head injection

* Make renderHead also yield an instruction

* Add changeset

* Add mdx test

.changeset/good-items-rest.md

@@ -0,0 +1,5 @@
+---
+'astro': patch
+---
+
+Ensure CSS injections properly when using multiple layouts

packages/astro/src/@types/astro.ts

@@ -1440,7 +1440,7 @@ export interface SSRResult {
 	links: Set<SSRElement>;
 	propagation: Map<string, PropagationHint>;
 	propagators: Map<AstroComponentFactory, AstroComponentInstance>;
-	extraHead: Array<any>;
+	extraHead: Array<string>;
 	cookies: AstroCookies | undefined;
 	createAstro(
 		Astro: AstroGlobalPartial,

packages/astro/src/runtime/server/render/astro/head-and-content.ts

@@ -4,7 +4,7 @@ const headAndContentSym = Symbol.for('astro.headAndContent');
 
 export type HeadAndContent = {
 	[headAndContentSym]: true;
-	head: string | RenderTemplateResult;
+	head: string;
 	content: RenderTemplateResult;
 };
 
@@ -13,7 +13,7 @@ export function isHeadAndContent(obj: unknown): obj is HeadAndContent {
 }
 
 export function createHeadAndContent(
-	head: string | RenderTemplateResult,
+	head: string,
 	content: RenderTemplateResult
 ): HeadAndContent {
 	return {

packages/astro/src/runtime/server/render/common.ts

@@ -2,6 +2,7 @@ import type { SSRResult } from '../../../@types/astro';
 import type { RenderInstruction } from './types.js';
 
 import { HTMLBytes, markHTMLString } from '../escape.js';
+import { renderAllHeadContent } from './head.js';
 import {
 	determineIfNeedsHydrationScript,
 	determinesIfNeedsDirectiveScript,
@@ -20,40 +21,48 @@ export const decoder = new TextDecoder();
 // These directive instructions bubble all the way up to renderPage so that we
 // can ensure they are added only once, and as soon as possible.
 export function stringifyChunk(result: SSRResult, chunk: string | SlotString | RenderInstruction) {
-	switch ((chunk as any).type) {
-		case 'directive': {
-			const { hydration } = chunk as RenderInstruction;
-			let needsHydrationScript = hydration && determineIfNeedsHydrationScript(result);
-			let needsDirectiveScript =
-				hydration && determinesIfNeedsDirectiveScript(result, hydration.directive);
-
-			let prescriptType: PrescriptType = needsHydrationScript
-				? 'both'
-				: needsDirectiveScript
-				? 'directive'
-				: null;
-			if (prescriptType) {
-				let prescripts = getPrescripts(prescriptType, hydration.directive);
-				return markHTMLString(prescripts);
-			} else {
-				return '';
-			}
-		}
-		default: {
-			if (isSlotString(chunk as string)) {
-				let out = '';
-				const c = chunk as SlotString;
-				if (c.instructions) {
-					for (const instr of c.instructions) {
-						out += stringifyChunk(result, instr);
-					}
+	if(typeof (chunk as any).type === 'string') {
+		const instruction = chunk as RenderInstruction;
+		switch(instruction.type) {
+			case 'directive': {
+				const { hydration } = instruction;
+				let needsHydrationScript = hydration && determineIfNeedsHydrationScript(result);
+				let needsDirectiveScript =
+					hydration && determinesIfNeedsDirectiveScript(result, hydration.directive);
+	
+				let prescriptType: PrescriptType = needsHydrationScript
+					? 'both'
+					: needsDirectiveScript
+					? 'directive'
+					: null;
+				if (prescriptType) {
+					let prescripts = getPrescripts(prescriptType, hydration.directive);
+					return markHTMLString(prescripts);
+				} else {
+					return '';
 				}
-				out += chunk.toString();
-				return out;
 			}
-
-			return chunk.toString();
+			case 'head': {
+				if(result._metadata.hasRenderedHead) {
+					return '';	
+				}
+				return renderAllHeadContent(result);
+			}
 		}
+	} else {
+		if (isSlotString(chunk as string)) {
+			let out = '';
+			const c = chunk as SlotString;
+			if (c.instructions) {
+				for (const instr of c.instructions) {
+					out += stringifyChunk(result, instr);
+				}
+			}
+			out += chunk.toString();
+			return out;
+		}
+
+		return chunk.toString();
 	}
 }
 

packages/astro/src/runtime/server/render/head.ts

@@ -1,7 +1,6 @@
 import type { SSRResult } from '../../../@types/astro';
 
 import { markHTMLString } from '../escape.js';
-import { renderChild } from './any.js';
 import { renderElement } from './util.js';
 
 // Filter out duplicate elements in our set
@@ -13,14 +12,8 @@ const uniqueElements = (item: any, index: number, all: any[]) => {
 	);
 };
 
-async function* renderExtraHead(result: SSRResult, base: string) {
-	yield base;
-	for (const part of result.extraHead) {
-		yield* renderChild(part);
-	}
-}
-
-function renderAllHeadContent(result: SSRResult) {
+export function renderAllHeadContent(result: SSRResult) {
+	result._metadata.hasRenderedHead = true;
 	const styles = Array.from(result.styles)
 		.filter(uniqueElements)
 		.map((style) => renderElement('style', style));
@@ -35,29 +28,31 @@ function renderAllHeadContent(result: SSRResult) {
 		.filter(uniqueElements)
 		.map((link) => renderElement('link', link, false));
 
-	const baseHeadContent = markHTMLString(links.join('\n') + styles.join('\n') + scripts.join('\n'));
+	let content = links.join('\n') + styles.join('\n') + scripts.join('\n');
 
 	if (result.extraHead.length > 0) {
-		return renderExtraHead(result, baseHeadContent);
-	} else {
-		return baseHeadContent;
+		for (const part of result.extraHead) {
+			content += part;
+		}
 	}
+
+	return markHTMLString(content);
 }
 
-export function createRenderHead(result: SSRResult) {
-	result._metadata.hasRenderedHead = true;
-	return renderAllHeadContent.bind(null, result);
+export function * renderHead(result: SSRResult) {
+	yield { type: 'head', result } as const;
 }
 
-export const renderHead = createRenderHead;
-
 // This function is called by Astro components that do not contain a <head> component
 // This accommodates the fact that using a <head> is optional in Astro, so this
 // is called before a component's first non-head HTML element. If the head was
 // already injected it is a noop.
-export async function* maybeRenderHead(result: SSRResult) {
+export function* maybeRenderHead(result: SSRResult) {
 	if (result._metadata.hasRenderedHead) {
 		return;
 	}
-	yield createRenderHead(result)();
+
+	// This is an instruction informing the page rendering that head might need rendering.
+	// This allows the page to deduplicate head injections.
+	yield { type: 'head', result } as const;
 }

packages/astro/src/runtime/server/render/slot.ts

@@ -26,7 +26,7 @@ export async function renderSlot(_result: any, slotted: string, fallback?: any):
 		let content = '';
 		let instructions: null | RenderInstruction[] = null;
 		for await (const chunk of iterator) {
-			if ((chunk as any).type === 'directive') {
+			if (typeof (chunk as any).type === 'string') {
 				if (instructions === null) {
 					instructions = [];
 				}

packages/astro/src/runtime/server/render/types.ts

@@ -1,8 +1,15 @@
 import type { SSRResult } from '../../../@types/astro';
 import type { HydrationMetadata } from '../hydration.js';
 
-export interface RenderInstruction {
+export type RenderDirectiveInstruction = {
 	type: 'directive';
 	result: SSRResult;
 	hydration: HydrationMetadata;
+};
+
+export type RenderHeadInstruction = {
+	type: 'head';
+	result: SSRResult;
 }
+
+export type RenderInstruction = RenderDirectiveInstruction | RenderHeadInstruction;

packages/astro/test/units/render/head.test.js

@@ -0,0 +1,181 @@
+import { expect } from 'chai';
+
+import {
+	createComponent,
+	render,
+	renderComponent,
+	renderSlot,
+	maybeRenderHead,
+	renderHead,
+	Fragment
+} from '../../../dist/runtime/server/index.js';
+import {
+	createBasicEnvironment,
+	createRenderContext,
+	renderPage,
+} from '../../../dist/core/render/index.js';
+import { defaultLogging as logging } from '../../test-utils.js';
+import * as cheerio from 'cheerio';
+
+const createAstroModule = (AstroComponent) => ({ default: AstroComponent });
+
+describe('core/render', () => {
+	describe('Injected head contents', () => {
+		let env;
+		before(async () => {
+			env = createBasicEnvironment({
+				logging,
+				renderers: [],
+			});
+		});
+
+		it('Multi-level layouts and head injection, with explicit head', async () => {
+			const BaseLayout = createComponent((result, _props, slots) => {
+				return render`<html>
+					<head>
+					${renderSlot(result, slots['head'])}
+					${renderHead(result)}
+					</head>
+					${maybeRenderHead(result)}
+					<body>
+						${renderSlot(result, slots['default'])}
+					</body>
+				</html>`;
+			})
+
+			const PageLayout = createComponent((result, _props, slots) => {
+				return render`${renderComponent(result, 'Layout', BaseLayout, {}, {
+						'default': () => render`
+							${maybeRenderHead(result)}
+							<main>
+								${renderSlot(result, slots['default'])}
+							</main>
+						`,
+						'head': () => render`
+							${renderComponent(result, 'Fragment', Fragment, { slot: 'head' }, {
+								'default': () => render`${renderSlot(result, slots['head'])}`
+							})}
+						`
+					})}
+				`;
+			});
+
+			const Page = createComponent((result, _props) => {
+				return render`${renderComponent(result, 'PageLayout', PageLayout, {}, {
+					'default': () => render`${maybeRenderHead(result)}<div>hello world</div>`,
+					'head': () => render`
+						${renderComponent(result, 'Fragment', Fragment, {slot: 'head'}, {
+							'default': () => render`<meta charset="utf-8">`
+						})}
+					`
+				})}`;
+			});
+
+			const ctx = createRenderContext({
+				request: new Request('http://example.com/'),
+				links: [
+					{ name: 'link', props: {rel:'stylesheet', href:'/main.css'}, children: '' }
+				]
+			});
+			const PageModule = createAstroModule(Page);
+
+			const response = await renderPage(PageModule, ctx, env);
+
+			const html = await response.text();
+			const $ = cheerio.load(html);
+
+			expect($('head link')).to.have.a.lengthOf(1);
+			expect($('body link')).to.have.a.lengthOf(0);
+		});
+
+		it('Multi-level layouts and head injection, without explicit head', async () => {
+			const BaseLayout = createComponent((result, _props, slots) => {
+				return render`<html>
+					${renderSlot(result, slots['head'])}
+					${maybeRenderHead(result)}
+					<body>
+						${renderSlot(result, slots['default'])}
+					</body>
+				</html>`;
+			})
+
+			const PageLayout = createComponent((result, _props, slots) => {
+				return render`${renderComponent(result, 'Layout', BaseLayout, {}, {
+						'default': () => render`
+							${maybeRenderHead(result)}
+							<main>
+								${renderSlot(result, slots['default'])}
+							</main>
+						`,
+						'head': () => render`
+							${renderComponent(result, 'Fragment', Fragment, { slot: 'head' }, {
+								'default': () => render`${renderSlot(result, slots['head'])}`
+							})}
+						`
+					})}
+				`;
+			});
+
+			const Page = createComponent((result, _props) => {
+				return render`${renderComponent(result, 'PageLayout', PageLayout, {}, {
+					'default': () => render`${maybeRenderHead(result)}<div>hello world</div>`,
+					'head': () => render`
+						${renderComponent(result, 'Fragment', Fragment, {slot: 'head'}, {
+							'default': () => render`<meta charset="utf-8">`
+						})}
+					`
+				})}`;
+			});
+
+			const ctx = createRenderContext({
+				request: new Request('http://example.com/'),
+				links: [
+					{ name: 'link', props: {rel:'stylesheet', href:'/main.css'}, children: '' }
+				]
+			});
+			const PageModule = createAstroModule(Page);
+
+			const response = await renderPage(PageModule, ctx, env);
+
+			const html = await response.text();
+			const $ = cheerio.load(html);
+
+			expect($('head link')).to.have.a.lengthOf(1);
+			expect($('body link')).to.have.a.lengthOf(0);
+		});
+
+		it('Multi-level layouts and head injection, without any content in layouts', async () => {
+			const BaseLayout = createComponent((result, _props, slots) => {
+				return render`${renderSlot(result, slots['default'])}`;
+			})
+
+			const PageLayout = createComponent((result, _props, slots) => {
+				return render`${renderComponent(result, 'Layout', BaseLayout, {}, {
+						'default': () => render`${renderSlot(result, slots['default'])}	`,
+					})}
+				`;
+			});
+
+			const Page = createComponent((result, _props) => {
+				return render`${renderComponent(result, 'PageLayout', PageLayout, {}, {
+					'default': () => render`${maybeRenderHead(result)}<div>hello world</div>`,
+				})}`;
+			});
+
+			const ctx = createRenderContext({
+				request: new Request('http://example.com/'),
+				links: [
+					{ name: 'link', props: {rel:'stylesheet', href:'/main.css'}, children: '' }
+				]
+			});
+			const PageModule = createAstroModule(Page);
+
+			const response = await renderPage(PageModule, ctx, env);
+
+			const html = await response.text();
+			const $ = cheerio.load(html);
+
+			expect($('link')).to.have.a.lengthOf(1);
+		});
+	});
+});

packages/integrations/mdx/test/css-head-mdx.test.js

@@ -0,0 +1,33 @@
+import mdx from '@astrojs/mdx';
+
+import { expect } from 'chai';
+import { parseHTML } from 'linkedom';
+import { loadFixture } from '../../../astro/test/test-utils.js';
+
+describe('Head injection w/ MDX', () => {
+	let fixture;
+
+	before(async () => {
+		fixture = await loadFixture({
+			root: new URL('./fixtures/css-head-mdx/', import.meta.url),
+			integrations: [mdx()],
+		});
+	});
+
+	describe('build', () => {
+		before(async () => {
+			await fixture.build();
+		});
+
+		it('only injects contents into head', async () => {
+			const html = await fixture.readFile('/indexThree/index.html');
+			const { document } = parseHTML(html);
+
+			const links = document.querySelectorAll('link[rel=stylesheet]');
+			expect(links).to.have.a.lengthOf(1);
+
+			const scripts = document.querySelectorAll('script[type=module]');
+			expect(scripts).to.have.a.lengthOf(1);
+		});
+	});
+});

packages/integrations/mdx/test/fixtures/css-head-mdx/src/components/HelloWorld.astro

@@ -0,0 +1,11 @@
+---
+---
+
+<h3>Hello world!!</h3>
+<slot />
+
+<style>h3 { color: red }</style>
+
+<script>
+console.log('hellooooo')
+</script>

packages/integrations/mdx/test/fixtures/css-head-mdx/src/layouts/One.astro

@@ -0,0 +1,15 @@
+---
+---
+
+<html lang="en">
+	<head>
+		<meta charset="utf-8" />
+		<link rel="icon" type="image/svg+xml" href="/favicon.svg" />
+		<meta name="viewport" content="width=device-width" />
+		<meta name="generator" content={Astro.generator} />
+		<title>Astro</title>
+	</head>
+	<body>
+		<slot />
+	</body>
+</html>

packages/integrations/mdx/test/fixtures/css-head-mdx/src/layouts/Three.astro

@@ -0,0 +1,6 @@
+---
+import Two from './Two.astro'
+---
+<Two>
+<slot />
+</Two>

packages/integrations/mdx/test/fixtures/css-head-mdx/src/layouts/Two.astro

@@ -0,0 +1,6 @@
+---
+import One from './One.astro'
+---
+<One>
+<slot />
+</One>

packages/integrations/mdx/test/fixtures/css-head-mdx/src/pages/indexOne.astro

@@ -0,0 +1,10 @@
+---
+import One from '../layouts/One.astro'
+
+import { Content } from '../test.mdx'
+---
+
+<One>
+	<h1>Astro</h1>
+	<Content />
+</One>

packages/integrations/mdx/test/fixtures/css-head-mdx/src/pages/indexThree.astro

@@ -0,0 +1,10 @@
+---
+import Three from '../layouts/Three.astro'
+
+import { Content } from '../test.mdx'
+---
+
+<Three>
+	<h1>Astro</h1>
+	<Content />
+</Three>

packages/integrations/mdx/test/fixtures/css-head-mdx/src/pages/indexTwo.astro

@@ -0,0 +1,10 @@
+---
+import Two from '../layouts/Two.astro'
+
+import { Content } from '../test.mdx'
+---
+
+<Two>
+	<h1>Astro</h1>
+	<Content />
+</Two>

packages/integrations/mdx/test/fixtures/css-head-mdx/src/pages/testOne.mdx

@@ -0,0 +1,15 @@
+---
+layout: '../layouts/One.astro'
+title: "hello world"
+publishDate: "2023-01-01"
+---
+
+import HelloWorld from '../components/HelloWorld.astro';
+
+# Test
+
+123
+
+<HelloWorld />
+
+456

packages/integrations/mdx/test/fixtures/css-head-mdx/src/pages/testThree.mdx

@@ -0,0 +1,15 @@
+---
+layout: '../layouts/Three.astro'
+title: "hello world"
+publishDate: "2023-01-01"
+---
+
+import HelloWorld from '../components/HelloWorld.astro';
+
+# Test
+
+123
+
+<HelloWorld />
+
+456

packages/integrations/mdx/test/fixtures/css-head-mdx/src/pages/testTwo.mdx

@@ -0,0 +1,15 @@
+---
+layout: '../layouts/Two.astro'
+title: "hello world"
+publishDate: "2023-01-01"
+---
+
+import HelloWorld from '../components/HelloWorld.astro';
+
+# Test
+
+123
+
+<HelloWorld />
+
+456

packages/integrations/mdx/test/fixtures/css-head-mdx/src/test.mdx

@@ -0,0 +1,14 @@
+---
+title: "hello world"
+publishDate: "2023-01-01"
+---
+
+import HelloWorld from './components/HelloWorld.astro';
+
+# Test
+
+123
+
+<HelloWorld />
+
+456