Refactor Astro page rendering (#7730)

6 files changed, 216 insertions, 141 deletions

diff --git a/packages/astro/src/runtime/server/jsx.ts b/packages/astro/src/runtime/server/jsx.ts
index 7afb70839..48f879b10 100644
--- a/packages/astro/src/runtime/server/jsx.ts
+++ b/packages/astro/src/runtime/server/jsx.ts
@@ -91,7 +91,11 @@ Did you forget to import the component or is it possible there is a typo?`);
 						props[key] = value;
 					}
 				}
-				const html = markHTMLString(await renderToString(result, vnode.type as any, props, slots));
+				const str = await renderToString(result, vnode.type as any, props, slots);
+				if (str instanceof Response) {
+					throw str;
+				}
+				const html = markHTMLString(str);
 				return html;
 			}
 			case !vnode.type && (vnode.type as any) !== 0:
diff --git a/packages/astro/src/runtime/server/render/astro/factory.ts b/packages/astro/src/runtime/server/render/astro/factory.ts
index 6d1b08563..97b8e4574 100644
--- a/packages/astro/src/runtime/server/render/astro/factory.ts
+++ b/packages/astro/src/runtime/server/render/astro/factory.ts
@@ -2,10 +2,6 @@ import type { PropagationHint, SSRResult } from '../../../../@types/astro';
 import type { HeadAndContent } from './head-and-content';
 import type { RenderTemplateResult } from './render-template';
 
-import { HTMLParts } from '../common.js';
-import { isHeadAndContent } from './head-and-content.js';
-import { renderAstroTemplateResult } from './render-template.js';
-
 export type AstroFactoryReturnValue = RenderTemplateResult | Response | HeadAndContent;
 
 // The callback passed to to $$createComponent
@@ -20,29 +16,6 @@ export function isAstroComponentFactory(obj: any): obj is AstroComponentFactory
 	return obj == null ? false : obj.isAstroComponentFactory === true;
 }
 
-// Calls a component and renders it into a string of HTML
-export async function renderToString(
-	result: SSRResult,
-	componentFactory: AstroComponentFactory,
-	props: any,
-	children: any
-): Promise<string> {
-	const factoryResult = await componentFactory(result, props, children);
-
-	if (factoryResult instanceof Response) {
-		const response = factoryResult;
-		throw response;
-	}
-
-	let parts = new HTMLParts();
-	const templateResult = isHeadAndContent(factoryResult) ? factoryResult.content : factoryResult;
-	for await (const chunk of renderAstroTemplateResult(templateResult)) {
-		parts.append(chunk, result);
-	}
-
-	return parts.toString();
-}
-
 export function isAPropagatingComponent(
 	result: SSRResult,
 	factory: AstroComponentFactory
diff --git a/packages/astro/src/runtime/server/render/astro/index.ts b/packages/astro/src/runtime/server/render/astro/index.ts
index cbddf7876..f7d9923ee 100644
--- a/packages/astro/src/runtime/server/render/astro/index.ts
+++ b/packages/astro/src/runtime/server/render/astro/index.ts
@@ -1,5 +1,5 @@
 export type { AstroComponentFactory } from './factory';
-export { isAstroComponentFactory, renderToString } from './factory.js';
+export { isAstroComponentFactory } from './factory.js';
 export { createHeadAndContent, isHeadAndContent } from './head-and-content.js';
 export type { AstroComponentInstance } from './instance';
 export { createAstroComponentInstance, isAstroComponentInstance } from './instance.js';
@@ -8,3 +8,4 @@ export {
 	renderAstroTemplateResult,
 	renderTemplate,
 } from './render-template.js';
+export { renderToReadableStream, renderToString } from './render.js';
diff --git a/packages/astro/src/runtime/server/render/astro/render.ts b/packages/astro/src/runtime/server/render/astro/render.ts
new file mode 100644
index 000000000..81b4375be
--- /dev/null
+++ b/packages/astro/src/runtime/server/render/astro/render.ts
@@ -0,0 +1,168 @@
+import type { RouteData, SSRResult } from '../../../../@types/astro';
+import { AstroError, AstroErrorData } from '../../../../core/errors/index.js';
+import { chunkToByteArray, chunkToString, encoder, type RenderDestination } from '../common.js';
+import type { AstroComponentFactory } from './factory.js';
+import { isHeadAndContent } from './head-and-content.js';
+import { isRenderTemplateResult, renderAstroTemplateResult } from './render-template.js';
+
+// Calls a component and renders it into a string of HTML
+export async function renderToString(
+	result: SSRResult,
+	componentFactory: AstroComponentFactory,
+	props: any,
+	children: any,
+	isPage = false,
+	route?: RouteData
+): Promise<string | Response> {
+	const templateResult = await callComponentAsTemplateResultOrResponse(
+		result,
+		componentFactory,
+		props,
+		children,
+		route
+	);
+
+	// If the Astro component returns a Response on init, return that response
+	if (templateResult instanceof Response) return templateResult;
+
+	let str = '';
+	let renderedFirstPageChunk = false;
+
+	const destination: RenderDestination = {
+		write(chunk) {
+			// Automatic doctype insertion for pages
+			if (isPage && !renderedFirstPageChunk) {
+				renderedFirstPageChunk = true;
+				if (!/<!doctype html/i.test(String(chunk))) {
+					const doctype = result.compressHTML ? '<!DOCTYPE html>' : '<!DOCTYPE html>\n';
+					str += doctype;
+				}
+			}
+
+			// `renderToString` doesn't work with emitting responses, so ignore here
+			if (chunk instanceof Response) return;
+
+			str += chunkToString(result, chunk);
+		},
+	};
+
+	for await (const chunk of renderAstroTemplateResult(templateResult)) {
+		destination.write(chunk);
+	}
+
+	return str;
+}
+
+// Calls a component and renders it into a readable stream
+export async function renderToReadableStream(
+	result: SSRResult,
+	componentFactory: AstroComponentFactory,
+	props: any,
+	children: any,
+	isPage = false,
+	route?: RouteData
+): Promise<ReadableStream | Response> {
+	const templateResult = await callComponentAsTemplateResultOrResponse(
+		result,
+		componentFactory,
+		props,
+		children,
+		route
+	);
+
+	// If the Astro component returns a Response on init, return that response
+	if (templateResult instanceof Response) return templateResult;
+
+	if (isPage) {
+		await bufferHeadContent(result);
+	}
+
+	let renderedFirstPageChunk = false;
+
+	return new ReadableStream({
+		start(controller) {
+			const destination: RenderDestination = {
+				write(chunk) {
+					// Automatic doctype insertion for pages
+					if (isPage && !renderedFirstPageChunk) {
+						renderedFirstPageChunk = true;
+						if (!/<!doctype html/i.test(String(chunk))) {
+							const doctype = result.compressHTML ? '<!DOCTYPE html>' : '<!DOCTYPE html>\n';
+							controller.enqueue(encoder.encode(doctype));
+						}
+					}
+
+					// `chunk` might be a Response that contains a redirect,
+					// that was rendered eagerly and therefore bypassed the early check
+					// whether headers can still be modified. In that case, throw an error
+					if (chunk instanceof Response) {
+						throw new AstroError({
+							...AstroErrorData.ResponseSentError,
+						});
+					}
+
+					const bytes = chunkToByteArray(result, chunk);
+					controller.enqueue(bytes);
+				},
+			};
+
+			(async () => {
+				try {
+					for await (const chunk of renderAstroTemplateResult(templateResult)) {
+						destination.write(chunk);
+					}
+					controller.close();
+				} catch (e) {
+					// We don't have a lot of information downstream, and upstream we can't catch the error properly
+					// So let's add the location here
+					if (AstroError.is(e) && !e.loc) {
+						e.setLocation({
+							file: route?.component,
+						});
+					}
+					controller.error(e);
+				}
+			})();
+		},
+	});
+}
+
+async function callComponentAsTemplateResultOrResponse(
+	result: SSRResult,
+	componentFactory: AstroComponentFactory,
+	props: any,
+	children: any,
+	route?: RouteData
+) {
+	const factoryResult = await componentFactory(result, props, children);
+
+	if (factoryResult instanceof Response) {
+		return factoryResult;
+	} else if (!isRenderTemplateResult(factoryResult)) {
+		throw new AstroError({
+			...AstroErrorData.OnlyResponseCanBeReturned,
+			message: AstroErrorData.OnlyResponseCanBeReturned.message(route?.route, typeof factoryResult),
+			location: {
+				file: route?.component,
+			},
+		});
+	}
+
+	return isHeadAndContent(factoryResult) ? factoryResult.content : factoryResult;
+}
+
+// Recursively calls component instances that might have head content
+// to be propagated up.
+async function bufferHeadContent(result: SSRResult) {
+	const iterator = result._metadata.propagators.values();
+	while (true) {
+		const { value, done } = iterator.next();
+		if (done) {
+			break;
+		}
+		const returnValue = await value.init(result);
+		if (isHeadAndContent(returnValue)) {
+			result._metadata.extraHead.push(returnValue.head);
+		}
+	}
+}
diff --git a/packages/astro/src/runtime/server/render/common.ts b/packages/astro/src/runtime/server/render/common.ts
index 50a99bc68..b2d41bd54 100644
--- a/packages/astro/src/runtime/server/render/common.ts
+++ b/packages/astro/src/runtime/server/render/common.ts
@@ -11,6 +11,14 @@ import {
 import { renderAllHeadContent } from './head.js';
 import { isSlotString, type SlotString } from './slot.js';
 
+export interface RenderDestination {
+	/**
+	 * Any rendering logic should call this to construct the HTML output.
+	 * See the `chunk` parameter for possible writable values
+	 */
+	write(chunk: string | HTMLBytes | RenderInstruction | Response): void;
+}
+
 export const Fragment = Symbol.for('astro:fragment');
 export const Renderer = Symbol.for('astro:renderer');
 
@@ -101,15 +109,22 @@ export class HTMLParts {
 	}
 }
 
+export function chunkToString(result: SSRResult, chunk: string | HTMLBytes | RenderInstruction) {
+	if (ArrayBuffer.isView(chunk)) {
+		return decoder.decode(chunk);
+	} else {
+		return stringifyChunk(result, chunk);
+	}
+}
+
 export function chunkToByteArray(
 	result: SSRResult,
 	chunk: string | HTMLBytes | RenderInstruction
 ): Uint8Array {
-	if (chunk instanceof Uint8Array) {
+	if (ArrayBuffer.isView(chunk)) {
 		return chunk as Uint8Array;
+	} else {
+		// stringify chunk might return a HTMLString
+		return encoder.encode(stringifyChunk(result, chunk));
 	}
-
-	// stringify chunk might return a HTMLString
-	let stringified = stringifyChunk(result, chunk);
-	return encoder.encode(stringified.toString());
 }
diff --git a/packages/astro/src/runtime/server/render/page.ts b/packages/astro/src/runtime/server/render/page.ts
index 4188d85e5..025b0b9e3 100644
--- a/packages/astro/src/runtime/server/render/page.ts
+++ b/packages/astro/src/runtime/server/render/page.ts
@@ -2,17 +2,12 @@ import type { RouteData, SSRResult } from '../../../@types/astro';
 import type { ComponentIterable } from './component';
 import type { AstroComponentFactory } from './index';
 
-import { AstroError, AstroErrorData } from '../../../core/errors/index.js';
+import { AstroError } from '../../../core/errors/index.js';
 import { isHTMLString } from '../escape.js';
 import { createResponse } from '../response.js';
-import {
-	isAstroComponentFactory,
-	isAstroComponentInstance,
-	isHeadAndContent,
-	isRenderTemplateResult,
-	renderAstroTemplateResult,
-} from './astro/index.js';
-import { HTMLParts, chunkToByteArray, encoder } from './common.js';
+import { isAstroComponentFactory, isAstroComponentInstance } from './astro/index.js';
+import { renderToReadableStream, renderToString } from './astro/render.js';
+import { HTMLParts, encoder } from './common.js';
 import { renderComponent } from './component.js';
 import { maybeRenderHead } from './head.js';
 
@@ -51,22 +46,6 @@ async function iterableToHTMLBytes(
 	return parts.toArrayBuffer();
 }
 
-// Recursively calls component instances that might have head content
-// to be propagated up.
-async function bufferHeadContent(result: SSRResult) {
-	const iterator = result._metadata.propagators.values();
-	while (true) {
-		const { value, done } = iterator.next();
-		if (done) {
-			break;
-		}
-		const returnValue = await value.init(result);
-		if (isHeadAndContent(returnValue)) {
-			result._metadata.extraHead.push(returnValue.head);
-		}
-	}
-}
-
 export async function renderPage(
 	result: SSRResult,
 	componentFactory: AstroComponentFactory | NonAstroPageComponent,
@@ -128,90 +107,25 @@ export async function renderPage(
 	// We avoid implicit head injection entirely.
 	result._metadata.headInTree =
 		result.componentMetadata.get(componentFactory.moduleId!)?.containsHead ?? false;
-	const factoryReturnValue = await componentFactory(result, props, children);
-	const factoryIsHeadAndContent = isHeadAndContent(factoryReturnValue);
-	if (isRenderTemplateResult(factoryReturnValue) || factoryIsHeadAndContent) {
-		// Wait for head content to be buffered up
-		await bufferHeadContent(result);
-		const templateResult = factoryIsHeadAndContent
-			? factoryReturnValue.content
-			: factoryReturnValue;
-
-		let iterable = renderAstroTemplateResult(templateResult);
-		let init = result.response;
-		let headers = new Headers(init.headers);
-		let body: BodyInit;
-
-		if (streaming) {
-			body = new ReadableStream({
-				start(controller) {
-					async function read() {
-						let i = 0;
-						try {
-							for await (const chunk of iterable) {
-								if (isHTMLString(chunk)) {
-									if (i === 0) {
-										if (!/<!doctype html/i.test(String(chunk))) {
-											controller.enqueue(
-												encoder.encode(
-													`${result.compressHTML ? '<!DOCTYPE html>' : '<!DOCTYPE html>\n'}`
-												)
-											);
-										}
-									}
-								}
-
-								// `chunk` might be a Response that contains a redirect,
-								// that was rendered eagerly and therefore bypassed the early check
-								// whether headers can still be modified. In that case, throw an error
-								if (chunk instanceof Response) {
-									throw new AstroError({
-										...AstroErrorData.ResponseSentError,
-									});
-								}
 
-								const bytes = chunkToByteArray(result, chunk);
-								controller.enqueue(bytes);
-								i++;
-							}
-							controller.close();
-						} catch (e) {
-							// We don't have a lot of information downstream, and upstream we can't catch the error properly
-							// So let's add the location here
-							if (AstroError.is(e) && !e.loc) {
-								e.setLocation({
-									file: route?.component,
-								});
-							}
-
-							controller.error(e);
-						}
-					}
-					read();
-				},
-			});
-		} else {
-			body = await iterableToHTMLBytes(result, iterable);
-			headers.set('Content-Length', body.byteLength.toString());
-		}
-
-		let response = createResponse(body, { ...init, headers });
-		return response;
+	let body: BodyInit | Response;
+	if (streaming) {
+		body = await renderToReadableStream(result, componentFactory, props, children, true, route);
+	} else {
+		body = await renderToString(result, componentFactory, props, children, true, route);
 	}
 
-	// We double check if the file return a Response
-	if (!(factoryReturnValue instanceof Response)) {
-		throw new AstroError({
-			...AstroErrorData.OnlyResponseCanBeReturned,
-			message: AstroErrorData.OnlyResponseCanBeReturned.message(
-				route?.route,
-				typeof factoryReturnValue
-			),
-			location: {
-				file: route?.component,
-			},
-		});
-	}
+	// If the Astro component returns a Response on init, return that response
+	if (body instanceof Response) return body;
 
-	return factoryReturnValue;
+	// Create final response from body
+	const init = result.response;
+	const headers = new Headers(init.headers);
+	// For non-streaming, convert string to byte array to calculate Content-Length
+	if (!streaming && typeof body === 'string') {
+		body = encoder.encode(body);
+		headers.set('Content-Length', body.byteLength.toString());
+	}
+	const response = createResponse(body, { ...init, headers });
+	return response;
 }