Prevent hydration scripts from being rendered in the wrong order (#4080)

* Prevent hydration scripts from being rendered in the wrong order

* Remove comment

* Update jsx

* Remove promise for stop

* Try skipping lit tests

* Stringify these chunks too

* Unskip lit

17 files changed, 199 insertions, 52 deletions

diff --git a/.changeset/ninety-news-change.md b/.changeset/ninety-news-change.md
new file mode 100644
index 000000000..8578678bf
--- /dev/null
+++ b/.changeset/ninety-news-change.md
@@ -0,0 +1,5 @@
+---
+'astro': patch
+---
+
+Fixes race condition for rendering hydration scripts
diff --git a/packages/astro/src/@types/astro.ts b/packages/astro/src/@types/astro.ts
index 570df260c..cf966a839 100644
--- a/packages/astro/src/@types/astro.ts
+++ b/packages/astro/src/@types/astro.ts
@@ -1092,6 +1092,8 @@ export interface SSRElement {
 export interface SSRMetadata {
 	renderers: SSRLoadedRenderer[];
 	pathname: string;
+	hasHydrationScript: boolean;
+	hasDirectives: Set<string>;
 }
 
 export interface SSRResult {
diff --git a/packages/astro/src/core/render/result.ts b/packages/astro/src/core/render/result.ts
index 744c32de6..32dcdd916 100644
--- a/packages/astro/src/core/render/result.ts
+++ b/packages/astro/src/core/render/result.ts
@@ -263,6 +263,8 @@ const canonicalURL = new URL(Astro.url.pathname, Astro.site);
 		_metadata: {
 			renderers,
 			pathname,
+			hasHydrationScript: false,
+			hasDirectives: new Set(),
 		},
 		response,
 	};
diff --git a/packages/astro/src/runtime/server/hydration.ts b/packages/astro/src/runtime/server/hydration.ts
index d58a72b4b..1e7760814 100644
--- a/packages/astro/src/runtime/server/hydration.ts
+++ b/packages/astro/src/runtime/server/hydration.ts
@@ -10,14 +10,16 @@ import { serializeListValue } from './util.js';
 
 const HydrationDirectives = ['load', 'idle', 'media', 'visible', 'only'];
 
+export interface HydrationMetadata {
+	directive: string;
+	value: string;
+	componentUrl: string;
+	componentExport: { value: string };
+};
+
 interface ExtractedProps {
 	isPage: boolean;
-	hydration: {
-		directive: string;
-		value: string;
-		componentUrl: string;
-		componentExport: { value: string };
-	} | null;
+	hydration: HydrationMetadata | null;
 	props: Record<string | number, any>;
 }
 
diff --git a/packages/astro/src/runtime/server/index.ts b/packages/astro/src/runtime/server/index.ts
index 9dbd4af41..d3c838266 100644
--- a/packages/astro/src/runtime/server/index.ts
+++ b/packages/astro/src/runtime/server/index.ts
@@ -10,7 +10,7 @@ import type {
 } from '../../@types/astro';
 
 import { escapeHTML, HTMLString, markHTMLString } from './escape.js';
-import { extractDirectives, generateHydrateScript } from './hydration.js';
+import { extractDirectives, generateHydrateScript, HydrationMetadata } from './hydration.js';
 import { createResponse } from './response.js';
 import {
 	determineIfNeedsHydrationScript,
@@ -130,12 +130,16 @@ export function createComponent(cb: AstroComponentFactory) {
 	return cb;
 }
 
-export async function renderSlot(_result: any, slotted: string, fallback?: any): Promise<string> {
+export async function renderSlot(result: any, slotted: string, fallback?: any): Promise<string> {
 	if (slotted) {
 		let iterator = _render(slotted);
 		let content = '';
 		for await (const chunk of iterator) {
-			content += chunk;
+			if((chunk as any).type === 'directive') {
+				content += stringifyChunk(result, chunk);
+			} else {
+				content += chunk;
+			}
 		}
 		return markHTMLString(content);
 	}
@@ -200,7 +204,7 @@ export async function renderComponent(
 	Component: unknown,
 	_props: Record<string | number, any>,
 	slots: any = {}
-): Promise<string | AsyncIterable<string>> {
+): Promise<string | AsyncIterable<string | RenderInstruction>> {
 	Component = await Component;
 	if (Component === Fragment) {
 		const children = await renderSlot(result, slots?.default);
@@ -226,7 +230,7 @@ export async function renderComponent(
 	}
 
 	if (Component && (Component as any).isAstroComponentFactory) {
-		async function* renderAstroComponentInline(): AsyncGenerator<string, void, undefined> {
+		async function* renderAstroComponentInline(): AsyncGenerator<string | RenderInstruction, void, undefined> {
 			let iterable = await renderToIterable(result, Component as any, _props, slots);
 			yield* iterable;
 		}
@@ -245,9 +249,6 @@ export async function renderComponent(
 
 	const { hydration, isPage, props } = extractDirectives(_props);
 	let html = '';
-	let needsHydrationScript = hydration && determineIfNeedsHydrationScript(result);
-	let needsDirectiveScript =
-		hydration && determinesIfNeedsDirectiveScript(result, hydration.directive);
 
 	if (hydration) {
 		metadata.hydrate = hydration.directive as AstroComponentMetadata['hydrate'];
@@ -481,15 +482,12 @@ If you're still stuck, please open an issue on GitHub or join us at https://astr
 		island.props['await-children'] = '';
 	}
 
-	// Scripts to prepend
-	let prescriptType: PrescriptType = needsHydrationScript
-		? 'both'
-		: needsDirectiveScript
-		? 'directive'
-		: null;
-	let prescripts = getPrescripts(prescriptType, hydration.directive);
+	async function * renderAll() {
+		yield { type: 'directive', hydration, result };
+		yield markHTMLString(renderElement('astro-island', island, false));
+	}
 
-	return markHTMLString(prescripts + renderElement('astro-island', island, false));
+	return renderAll();
 }
 
 /** Create the Astro.fetchContent() runtime function. */
@@ -758,7 +756,7 @@ export async function renderToIterable(
 	componentFactory: AstroComponentFactory,
 	props: any,
 	children: any
-): Promise<AsyncIterable<string>> {
+): Promise<AsyncIterable<string | RenderInstruction>> {
 	const Component = await componentFactory(result, props, children);
 
 	if (!isAstroComponent(Component)) {
@@ -775,6 +773,35 @@ export async function renderToIterable(
 
 const encoder = new TextEncoder();
 
+// 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) {
+	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: {
+			return chunk.toString();
+		}
+	}
+}
+
 export async function renderPage(
 	result: SSRResult,
 	componentFactory: AstroComponentFactory,
@@ -814,6 +841,7 @@ export async function renderPage(
 		let init = result.response;
 		let headers = new Headers(init.headers);
 		let body: BodyInit;
+
 		if (streaming) {
 			body = new ReadableStream({
 				start(controller) {
@@ -821,7 +849,8 @@ export async function renderPage(
 						let i = 0;
 						try {
 							for await (const chunk of iterable) {
-								let html = chunk.toString();
+								let html = stringifyChunk(result, chunk);
+
 								if (i === 0) {
 									if (!/<!doctype html/i.test(html)) {
 										controller.enqueue(encoder.encode('<!DOCTYPE html>\n'));
@@ -842,13 +871,13 @@ export async function renderPage(
 			body = '';
 			let i = 0;
 			for await (const chunk of iterable) {
-				let html = chunk.toString();
+				let html = stringifyChunk(result, chunk);
 				if (i === 0) {
 					if (!/<!doctype html/i.test(html)) {
 						body += '<!DOCTYPE html>\n';
 					}
 				}
-				body += chunk;
+				body += html;
 				i++;
 			}
 			const bytes = encoder.encode(body);
@@ -901,13 +930,28 @@ export async function* maybeRenderHead(result: SSRResult): AsyncIterable<string>
 	yield renderHead(result);
 }
 
+export interface RenderInstruction {
+	type: 'directive';
+	result: SSRResult;
+	hydration: HydrationMetadata;
+}
+
 export async function* renderAstroComponent(
 	component: InstanceType<typeof AstroComponent>
-): AsyncIterable<string> {
+): AsyncIterable<string | RenderInstruction> {
 	for await (const value of component) {
 		if (value || value === 0) {
 			for await (const chunk of _render(value)) {
-				yield markHTMLString(chunk);
+				switch(chunk.type) {
+					case 'directive': {
+						yield chunk;
+						break;
+					}
+					default: {
+						yield markHTMLString(chunk);
+						break;
+					}
+				}
 			}
 		}
 	}
diff --git a/packages/astro/src/runtime/server/jsx.ts b/packages/astro/src/runtime/server/jsx.ts
index 687d0c9b9..add4c0ca4 100644
--- a/packages/astro/src/runtime/server/jsx.ts
+++ b/packages/astro/src/runtime/server/jsx.ts
@@ -6,9 +6,11 @@ import {
 	escapeHTML,
 	HTMLString,
 	markHTMLString,
+	RenderInstruction,
 	renderComponent,
 	renderToString,
 	spreadAttributes,
+	stringifyChunk,
 	voidElementNames,
 } from './index.js';
 
@@ -119,7 +121,7 @@ export async function renderJSX(result: SSRResult, vnode: any): Promise<any> {
 			}
 			await Promise.all(slotPromises);
 
-			let output: string | AsyncIterable<string>;
+			let output: string | AsyncIterable<string | RenderInstruction>;
 			if (vnode.type === ClientOnlyPlaceholder && vnode.props['client:only']) {
 				output = await renderComponent(
 					result,
@@ -137,7 +139,16 @@ export async function renderJSX(result: SSRResult, vnode: any): Promise<any> {
 					slots
 				);
 			}
-			return markHTMLString(output);
+			if(typeof output !== 'string' && Symbol.asyncIterator in output) {
+				let body = '';
+				for await (const chunk of output) {
+					let html = stringifyChunk(result, chunk);
+					body += html;
+				}
+				return markHTMLString(body);
+			} else {
+				return markHTMLString(output);
+			}
 		}
 	}
 	// numbers, plain objects, etc
diff --git a/packages/astro/src/runtime/server/scripts.ts b/packages/astro/src/runtime/server/scripts.ts
index 2dba232f9..ab073fe67 100644
--- a/packages/astro/src/runtime/server/scripts.ts
+++ b/packages/astro/src/runtime/server/scripts.ts
@@ -7,17 +7,11 @@ import onlyPrebuilt from '../client/only.prebuilt.js';
 import visiblePrebuilt from '../client/visible.prebuilt.js';
 import islandScript from './astro-island.prebuilt.js';
 
-// This is used to keep track of which requests (pages) have had the hydration script
-// appended. We only add the hydration script once per page, and since the SSRResult
-// object corresponds to one page request, we are using it as a key to know.
-const resultsWithHydrationScript = new WeakSet<SSRResult>();
-
 export function determineIfNeedsHydrationScript(result: SSRResult): boolean {
-	if (resultsWithHydrationScript.has(result)) {
+	if(result._metadata.hasHydrationScript) {
 		return false;
 	}
-	resultsWithHydrationScript.add(result);
-	return true;
+	return result._metadata.hasHydrationScript = true;
 }
 
 export const hydrationScripts: Record<string, string> = {
@@ -28,17 +22,11 @@ export const hydrationScripts: Record<string, string> = {
 	visible: visiblePrebuilt,
 };
 
-const resultsWithDirectiveScript = new Map<string, WeakSet<SSRResult>>();
-
 export function determinesIfNeedsDirectiveScript(result: SSRResult, directive: string): boolean {
-	if (!resultsWithDirectiveScript.has(directive)) {
-		resultsWithDirectiveScript.set(directive, new WeakSet());
-	}
-	const set = resultsWithDirectiveScript.get(directive)!;
-	if (set.has(result)) {
+	if(result._metadata.hasDirectives.has(directive)) {
 		return false;
 	}
-	set.add(result);
+	result._metadata.hasDirectives.add(directive);
 	return true;
 }
 
diff --git a/packages/astro/test/fixtures/hydration-race/astro.config.mjs b/packages/astro/test/fixtures/hydration-race/astro.config.mjs
new file mode 100644
index 000000000..c9662ed09
--- /dev/null
+++ b/packages/astro/test/fixtures/hydration-race/astro.config.mjs
@@ -0,0 +1,5 @@
+import preact from '@astrojs/preact';
+
+export default {
+	integrations: [preact()]
+};
diff --git a/packages/astro/test/fixtures/hydration-race/package.json b/packages/astro/test/fixtures/hydration-race/package.json
new file mode 100644
index 000000000..57882c32b
--- /dev/null
+++ b/packages/astro/test/fixtures/hydration-race/package.json
@@ -0,0 +1,13 @@
+{
+  "name": "@test/hydration-race",
+  "version": "0.0.0",
+  "private": true,
+  "scripts": {
+    "dev": "astro dev",
+    "build": "astro build"
+  },
+  "dependencies": {
+    "astro": "workspace:*",
+    "@astrojs/preact": "workspace:*"
+  }
+}
diff --git a/packages/astro/test/fixtures/hydration-race/src/components/Deeper.astro b/packages/astro/test/fixtures/hydration-race/src/components/Deeper.astro
new file mode 100644
index 000000000..76945dff9
--- /dev/null
+++ b/packages/astro/test/fixtures/hydration-race/src/components/Deeper.astro
@@ -0,0 +1,9 @@
+---
+import One from './One.jsx';
+---
+
+<div>
+	<span>Before one</span>
+	<One name="One" client:idle />
+</div>
+
diff --git a/packages/astro/test/fixtures/hydration-race/src/components/One.jsx b/packages/astro/test/fixtures/hydration-race/src/components/One.jsx
new file mode 100644
index 000000000..35dd7753a
--- /dev/null
+++ b/packages/astro/test/fixtures/hydration-race/src/components/One.jsx
@@ -0,0 +1,7 @@
+import { h } from 'preact';
+
+export default function({ name }) {
+	return (
+		<div>Hello {name} </div>
+	);
+}
diff --git a/packages/astro/test/fixtures/hydration-race/src/components/Wrapper.astro b/packages/astro/test/fixtures/hydration-race/src/components/Wrapper.astro
new file mode 100644
index 000000000..c751b2f77
--- /dev/null
+++ b/packages/astro/test/fixtures/hydration-race/src/components/Wrapper.astro
@@ -0,0 +1,11 @@
+---
+import One from './One.jsx';
+import Deeper from './Deeper.astro';
+---
+
+<div>
+	<Deeper />
+	<span>Before one</span>
+	<One name="Two" client:idle />
+</div>
+
diff --git a/packages/astro/test/fixtures/hydration-race/src/pages/index.astro b/packages/astro/test/fixtures/hydration-race/src/pages/index.astro
new file mode 100644
index 000000000..d5272cc75
--- /dev/null
+++ b/packages/astro/test/fixtures/hydration-race/src/pages/index.astro
@@ -0,0 +1,16 @@
+---
+import Wrapper from '../components/Wrapper.astro';
+import One from '../components/One.jsx';
+---
+<html>
+<head>
+	<title>Testing</title>
+</head>
+<body>
+	<main>
+		<Wrapper />
+	</main>
+
+	<One name="Three" client:idle />
+</body>
+</html>
diff --git a/packages/astro/test/hydration-race.test.js b/packages/astro/test/hydration-race.test.js
new file mode 100644
index 000000000..d35681557
--- /dev/null
+++ b/packages/astro/test/hydration-race.test.js
@@ -0,0 +1,29 @@
+import { expect } from 'chai';
+import * as cheerio from 'cheerio';
+import { loadFixture } from './test-utils.js';
+
+describe('Hydration script ordering', async () => {
+	let fixture;
+
+	before(async () => {
+		fixture = await loadFixture({ root: './fixtures/hydration-race' });
+		await fixture.build();
+	});
+
+	it('Places the hydration script before the first island', async () => {
+		let html = await fixture.readFile('/index.html');
+		let $ = cheerio.load(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 prevSibling = firstIsland.prev();
+		expect(prevSibling.prop('tagName')).to.equal('SCRIPT');
+
+		// 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');
+	});
+});
diff --git a/packages/integrations/cloudflare/test/basics.test.js b/packages/integrations/cloudflare/test/basics.test.js
index e6fe6642e..7e3e9bb93 100644
--- a/packages/integrations/cloudflare/test/basics.test.js
+++ b/packages/integrations/cloudflare/test/basics.test.js
@@ -25,7 +25,7 @@ describe('Basic app', () => {
 			let $ = cheerio.load(html);
 			expect($('h1').text()).to.equal('Testing');
 		} finally {
-			await stop();
+			stop();
 		}
 	});
 });
diff --git a/packages/integrations/cloudflare/test/test-utils.js b/packages/integrations/cloudflare/test/test-utils.js
index e0f521949..58cb8f9dd 100644
--- a/packages/integrations/cloudflare/test/test-utils.js
+++ b/packages/integrations/cloudflare/test/test-utils.js
@@ -57,11 +57,6 @@ export function runCLI(basePath, { silent }) {
 		ready,
 		stop() {
 			p.kill();
-			return new Promise((resolve) => {
-				p.addListener('exit', () => {
-					resolve();
-				});
-			});
 		},
 	};
 }
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
index 447cc045c..2263b4618 100644
--- a/pnpm-lock.yaml
+++ b/pnpm-lock.yaml
@@ -1556,6 +1556,14 @@ importers:
     dependencies:
       astro: link:../../..
 
+  packages/astro/test/fixtures/hydration-race:
+    specifiers:
+      '@astrojs/preact': workspace:*
+      astro: workspace:*
+    dependencies:
+      '@astrojs/preact': link:../../../../integrations/preact
+      astro: link:../../..
+
   packages/astro/test/fixtures/import-ts-with-js:
     specifiers:
       astro: workspace:*