Ensure hydration scripts inside of slots render ASAP (#4288)

* Ensure hydration scripts inside of slots render ASAP

* Changeset

* fix type errors

* Update packages/astro/src/runtime/server/render/page.ts

Co-authored-by: Nate Moore <natemoo-re@users.noreply.github.com>

Co-authored-by: Nate Moore <natemoo-re@users.noreply.github.com>

16 files changed, 111 insertions, 35 deletions

diff --git a/.changeset/four-students-rescue.md b/.changeset/four-students-rescue.md
new file mode 100644
index 000000000..c2939acd2
--- /dev/null
+++ b/.changeset/four-students-rescue.md
@@ -0,0 +1,5 @@
+---
+'astro': patch
+---
+
+Ensure hydration scripts inside of slots render ASAP
diff --git a/packages/astro/src/@types/astro.ts b/packages/astro/src/@types/astro.ts
index abdfe3dea..b45b01b5b 100644
--- a/packages/astro/src/@types/astro.ts
+++ b/packages/astro/src/@types/astro.ts
@@ -1112,11 +1112,25 @@ export interface SSRElement {
 	children: string;
 }
 
+export interface HydrationMetadata {
+	directive: string;
+	value: string;
+	componentUrl: string;
+	componentExport: { value: string };
+}
+
+export interface SSRRenderInstruction {
+	type: 'directive';
+	result: SSRResult;
+	hydration: HydrationMetadata;
+}
+
 export interface SSRMetadata {
 	renderers: SSRLoadedRenderer[];
 	pathname: string;
 	hasHydrationScript: boolean;
 	hasDirectives: Set<string>;
+	pendingInstructions: Set<SSRRenderInstruction>;
 }
 
 export interface SSRResult {
diff --git a/packages/astro/src/core/render/result.ts b/packages/astro/src/core/render/result.ts
index 32dcdd916..37b07cb06 100644
--- a/packages/astro/src/core/render/result.ts
+++ b/packages/astro/src/core/render/result.ts
@@ -265,6 +265,7 @@ const canonicalURL = new URL(Astro.url.pathname, Astro.site);
 			pathname,
 			hasHydrationScript: false,
 			hasDirectives: new Set(),
+			pendingInstructions: new Set(),
 		},
 		response,
 	};
diff --git a/packages/astro/src/runtime/server/hydration.ts b/packages/astro/src/runtime/server/hydration.ts
index f100c1d7c..3f0ad256e 100644
--- a/packages/astro/src/runtime/server/hydration.ts
+++ b/packages/astro/src/runtime/server/hydration.ts
@@ -1,5 +1,6 @@
 import type {
 	AstroComponentMetadata,
+	HydrationMetadata,
 	SSRElement,
 	SSRLoadedRenderer,
 	SSRResult,
@@ -12,12 +13,7 @@ const HydrationDirectivesRaw = ['load', 'idle', 'media', 'visible', 'only'];
 const HydrationDirectives = new Set(HydrationDirectivesRaw);
 export const HydrationDirectiveProps = new Set(HydrationDirectivesRaw.map((n) => `client:${n}`));
 
-export interface HydrationMetadata {
-	directive: string;
-	value: string;
-	componentUrl: string;
-	componentExport: { value: string };
-}
+export { HydrationMetadata };
 
 interface ExtractedProps {
 	isPage: boolean;
diff --git a/packages/astro/src/runtime/server/index.ts b/packages/astro/src/runtime/server/index.ts
index 1304aca3e..c80811843 100644
--- a/packages/astro/src/runtime/server/index.ts
+++ b/packages/astro/src/runtime/server/index.ts
@@ -26,7 +26,7 @@ export {
 	stringifyChunk,
 	voidElementNames,
 } from './render/index.js';
-export type { AstroComponentFactory, RenderInstruction } from './render/index.js';
+export type { AstroComponentFactory } from './render/index.js';
 import type { AstroComponentFactory } from './render/index.js';
 
 import { markHTMLString } from './escape.js';
diff --git a/packages/astro/src/runtime/server/jsx.ts b/packages/astro/src/runtime/server/jsx.ts
index 6041d514d..d38a57184 100644
--- a/packages/astro/src/runtime/server/jsx.ts
+++ b/packages/astro/src/runtime/server/jsx.ts
@@ -1,12 +1,11 @@
 /* eslint-disable no-console */
-import { SSRResult } from '../../@types/astro.js';
+import { SSRResult, SSRRenderInstruction } from '../../@types/astro.js';
 import { AstroJSX, isVNode } from '../../jsx-runtime/index.js';
 import {
 	escapeHTML,
 	HTMLString,
 	markHTMLString,
 	renderComponent,
-	RenderInstruction,
 	renderToString,
 	spreadAttributes,
 	stringifyChunk,
@@ -122,7 +121,7 @@ export async function renderJSX(result: SSRResult, vnode: any): Promise<any> {
 			}
 			await Promise.all(slotPromises);
 
-			let output: string | AsyncIterable<string | RenderInstruction>;
+			let output: string | AsyncIterable<string | SSRRenderInstruction>;
 			if (vnode.type === ClientOnlyPlaceholder && vnode.props['client:only']) {
 				output = await renderComponent(
 					result,
diff --git a/packages/astro/src/runtime/server/render/any.ts b/packages/astro/src/runtime/server/render/any.ts
index 9c6e199c2..ecb26973e 100644
--- a/packages/astro/src/runtime/server/render/any.ts
+++ b/packages/astro/src/runtime/server/render/any.ts
@@ -1,3 +1,4 @@
+import { SSRRenderInstruction, SSRResult } from '../../../@types/astro';
 import { escapeHTML, HTMLString, markHTMLString } from '../escape.js';
 import { AstroComponent, renderAstroComponent } from './astro.js';
 import { stringifyChunk } from './common.js';
@@ -34,13 +35,15 @@ export async function* renderChild(child: any): AsyncIterable<any> {
 	}
 }
 
-export async function renderSlot(result: any, slotted: string, fallback?: any): Promise<string> {
+export async function renderSlot(result: SSRResult, slotted: string, fallback?: any): Promise<string> {
 	if (slotted) {
 		let iterator = renderChild(slotted);
 		let content = '';
 		for await (const chunk of iterator) {
 			if ((chunk as any).type === 'directive') {
-				content += stringifyChunk(result, chunk);
+				// Do not render directives scripts in a slot, instead mark them as pending
+				// so that renderPage includes them ASAP.
+				result._metadata.pendingInstructions.add(chunk as SSRRenderInstruction);
 			} else {
 				content += chunk;
 			}
diff --git a/packages/astro/src/runtime/server/render/astro.ts b/packages/astro/src/runtime/server/render/astro.ts
index b7e47a68c..94229449e 100644
--- a/packages/astro/src/runtime/server/render/astro.ts
+++ b/packages/astro/src/runtime/server/render/astro.ts
@@ -1,6 +1,5 @@
-import type { SSRResult } from '../../../@types/astro';
+import type { SSRResult, SSRRenderInstruction } from '../../../@types/astro';
 import type { AstroComponentFactory } from './index';
-import type { RenderInstruction } from './types';
 
 import { markHTMLString } from '../escape.js';
 import { HydrationDirectiveProps } from '../hydration.js';
@@ -58,7 +57,7 @@ export function isAstroComponent(obj: any): obj is AstroComponent {
 
 export async function* renderAstroComponent(
 	component: InstanceType<typeof AstroComponent>
-): AsyncIterable<string | RenderInstruction> {
+): AsyncIterable<string | SSRRenderInstruction> {
 	for await (const value of component) {
 		if (value || value === 0) {
 			for await (const chunk of renderChild(value)) {
@@ -104,7 +103,7 @@ export async function renderToIterable(
 	displayName: string,
 	props: any,
 	children: any
-): Promise<AsyncIterable<string | RenderInstruction>> {
+): Promise<AsyncIterable<string | SSRRenderInstruction>> {
 	validateComponentProps(props, displayName);
 	const Component = await componentFactory(result, props, children);
 
diff --git a/packages/astro/src/runtime/server/render/common.ts b/packages/astro/src/runtime/server/render/common.ts
index 27b012aa4..db1f8e8d3 100644
--- a/packages/astro/src/runtime/server/render/common.ts
+++ b/packages/astro/src/runtime/server/render/common.ts
@@ -1,5 +1,4 @@
-import type { SSRResult } from '../../../@types/astro';
-import type { RenderInstruction } from './types.js';
+import type { SSRResult, SSRRenderInstruction } from '../../../@types/astro';
 
 import { markHTMLString } from '../escape.js';
 import {
@@ -15,10 +14,10 @@ export const Renderer = Symbol.for('astro:renderer');
 // Rendering produces either marked strings of HTML or instructions for hydration.
 // 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 | RenderInstruction) {
+export function stringifyChunk(result: SSRResult, chunk: string | SSRRenderInstruction) {
 	switch ((chunk as any).type) {
 		case 'directive': {
-			const { hydration } = chunk as RenderInstruction;
+			const { hydration } = chunk as SSRRenderInstruction;
 			let needsHydrationScript = hydration && determineIfNeedsHydrationScript(result);
 			let needsDirectiveScript =
 				hydration && determinesIfNeedsDirectiveScript(result, hydration.directive);
diff --git a/packages/astro/src/runtime/server/render/component.ts b/packages/astro/src/runtime/server/render/component.ts
index 18dc5a00d..43215d0cc 100644
--- a/packages/astro/src/runtime/server/render/component.ts
+++ b/packages/astro/src/runtime/server/render/component.ts
@@ -1,5 +1,5 @@
-import type { AstroComponentMetadata, SSRLoadedRenderer, SSRResult } from '../../../@types/astro';
-import type { RenderInstruction } from './types.js';
+import type { AstroComponentMetadata, SSRLoadedRenderer, SSRResult, SSRRenderInstruction } from '../../../@types/astro';
+
 
 import { markHTMLString } from '../escape.js';
 import { extractDirectives, generateHydrateScript } from '../hydration.js';
@@ -49,7 +49,7 @@ export async function renderComponent(
 	Component: unknown,
 	_props: Record<string | number, any>,
 	slots: any = {}
-): Promise<string | AsyncIterable<string | RenderInstruction>> {
+): Promise<string | AsyncIterable<string | SSRRenderInstruction>> {
 	Component = await Component;
 
 	switch (getComponentType(Component)) {
@@ -79,7 +79,7 @@ export async function renderComponent(
 
 		case 'astro-factory': {
 			async function* renderAstroComponentInline(): AsyncGenerator<
-				string | RenderInstruction,
+				string | SSRRenderInstruction,
 				void,
 				undefined
 			> {
diff --git a/packages/astro/src/runtime/server/render/index.ts b/packages/astro/src/runtime/server/render/index.ts
index a7de515e2..1567307d5 100644
--- a/packages/astro/src/runtime/server/render/index.ts
+++ b/packages/astro/src/runtime/server/render/index.ts
@@ -7,7 +7,6 @@ export { renderComponent } from './component.js';
 export { renderHTMLElement } from './dom.js';
 export { maybeRenderHead, renderHead } from './head.js';
 export { renderPage } from './page.js';
-export type { RenderInstruction } from './types';
 export { addAttribute, defineScriptVars, voidElementNames } from './util.js';
 
 // The callback passed to to $$createComponent
diff --git a/packages/astro/src/runtime/server/render/page.ts b/packages/astro/src/runtime/server/render/page.ts
index 9e2c0fd2b..2a372df67 100644
--- a/packages/astro/src/runtime/server/render/page.ts
+++ b/packages/astro/src/runtime/server/render/page.ts
@@ -49,13 +49,28 @@ export async function renderPage(
 		let headers = new Headers(init.headers);
 		let body: BodyInit;
 
+		// Combines HTML chunks coming from the iterable with rendering instructions
+		// added to metadata. These instructions need to go out first to ensure
+		// the scripts exist before the islands that need them.
+		async function * eachChunk() {
+			for await (const chunk of iterable) {
+				if(result._metadata.pendingInstructions.size) {
+					for(const instr of result._metadata.pendingInstructions) {
+						result._metadata.pendingInstructions.delete(instr);
+						yield instr;
+					}
+				}
+				yield chunk;
+			}
+		}
+
 		if (streaming) {
 			body = new ReadableStream({
 				start(controller) {
 					async function read() {
 						let i = 0;
 						try {
-							for await (const chunk of iterable) {
+							for await (const chunk of eachChunk()) {
 								let html = stringifyChunk(result, chunk);
 
 								if (i === 0) {
@@ -77,7 +92,7 @@ export async function renderPage(
 		} else {
 			body = '';
 			let i = 0;
-			for await (const chunk of iterable) {
+			for await (const chunk of eachChunk()) {
 				let html = stringifyChunk(result, chunk);
 				if (i === 0) {
 					if (!/<!doctype html/i.test(html)) {
diff --git a/packages/astro/src/runtime/server/render/types.ts b/packages/astro/src/runtime/server/render/types.ts
deleted file mode 100644
index a9e1afe65..000000000
--- a/packages/astro/src/runtime/server/render/types.ts
+++ /dev/null
@@ -1,8 +0,0 @@
-import type { SSRResult } from '../../../@types/astro';
-import type { HydrationMetadata } from '../hydration.js';
-
-export interface RenderInstruction {
-	type: 'directive';
-	result: SSRResult;
-	hydration: HydrationMetadata;
-}
diff --git a/packages/astro/test/fixtures/hydration-race/src/components/Layout.astro b/packages/astro/test/fixtures/hydration-race/src/components/Layout.astro
new file mode 100644
index 000000000..8ee14c08f
--- /dev/null
+++ b/packages/astro/test/fixtures/hydration-race/src/components/Layout.astro
@@ -0,0 +1,6 @@
+---
+import Wrapper from './Wrapper.astro';
+---
+
+<Wrapper />
+<slot />
diff --git a/packages/astro/test/fixtures/hydration-race/src/pages/slot.astro b/packages/astro/test/fixtures/hydration-race/src/pages/slot.astro
new file mode 100644
index 000000000..7f054d45f
--- /dev/null
+++ b/packages/astro/test/fixtures/hydration-race/src/pages/slot.astro
@@ -0,0 +1,15 @@
+---
+import Layout from '../components/Layout.astro';
+import Wrapper from '../components/Wrapper.astro';
+import One from '../components/One.jsx';
+---
+<html>
+<head>
+	<title>Testing</title>
+</head>
+<body>
+	<Layout>
+		<One name="Three" client:idle />
+	</Layout>
+</body>
+</html>
diff --git a/packages/astro/test/hydration-race.test.js b/packages/astro/test/hydration-race.test.js
index d35681557..c26fa74fb 100644
--- a/packages/astro/test/hydration-race.test.js
+++ b/packages/astro/test/hydration-race.test.js
@@ -26,4 +26,37 @@ describe('Hydration script ordering', async () => {
 		expect($('style')).to.have.a.lengthOf(1, 'hydration style added once');
 		expect($('script')).to.have.a.lengthOf(1, 'only one hydration script needed');
 	});
+
+	it('Hydration placed in correct location when there is slots used', async () => {
+		let html = await fixture.readFile('/slot/index.html');
+		let $ = cheerio.load(html);
+		//console.log(html);
+
+		// First, let's make sure all islands rendered (or test is bad)
+		expect($('astro-island')).to.have.a.lengthOf(3);
+
+		// Now let's make sure the hydration script is placed before the first component
+		let firstIsland = $($('astro-island').get(0));
+		let el = firstIsland.prev();
+		let foundScript = false;
+
+		// Traverse the DOM going backwards until we find a script, if there is one.
+		while(el.length) {
+			let last = el;
+			while(el.length) {
+				if(el.prop('tagName') === 'SCRIPT') {
+					foundScript = true;
+				}
+				last = el;
+				el = el.prev();
+			}
+			el = last.parent();
+		}
+
+		expect(foundScript).to.be.true;
+
+		// Sanity check that we're only rendering them once.
+		expect($('style')).to.have.a.lengthOf(1, 'hydration style added once');
+		expect($('script')).to.have.a.lengthOf(1, 'only one hydration script needed');
+	});
 });