feat: Inject state at end of body

packages/request-state/package.json

@@ -33,8 +33,7 @@
     "@inox-tools/utils": "workspace:^",
     "astro-integration-kit": "catalog:",
     "content-type": "catalog:",
-    "devalue": "catalog:",
-    "parse5-html-rewriting-stream": "catalog:"
+    "devalue": "catalog:"
   },
   "devDependencies": {
     "@types/content-type": "catalog:",

packages/request-state/src/runtime/middleware.ts

@@ -1,7 +1,6 @@
 import { defineMiddleware } from 'astro/middleware';
 import { collectState } from './serverState.js';
 import { parse } from 'content-type';
-import { RewritingStream } from 'parse5-html-rewriting-stream';
 
 export const onRequest = defineMiddleware(async (_, next) => {
 	const { getState, result } = await collectState(next);
@@ -14,56 +13,31 @@ export const onRequest = defineMiddleware(async (_, next) => {
 
 	if (mediaType !== 'text/html' && !mediaType.startsWith('text/html+')) return result;
 
-	async function render() {
-		const rewriter = new RewritingStream();
+	const newBody = result.body
+		?.pipeThrough(new TextDecoderStream())
+		.pipeThrough(injectState(getState))
+		.pipeThrough(new TextEncoderStream());
 
-		let injected = false;
+	return new Response(newBody, result);
+});
 
-		rewriter.on('endTag', (tag) => {
-			if (!injected && tag.tagName === 'html') {
-				const state = getState();
-				if (state) {
-					rewriter.emitRaw(
-						`<script id="it-astro-state" type="application/json+devalue">${state}</script>`
-					);
+function injectState(getState: () => string | false) {
+	let injected = false;
+	return new TransformStream({
+		transform(chunk, controller) {
+			if (!injected) {
+				const bodyCloseIndex = chunk.indexOf('</body>');
+				if (bodyCloseIndex > -1) {
+					const state = getState();
+					if (state) {
+						const stateScript = `<script id="it-astro-state" type="application/json+devalue">${state}</script>`;
+
+						chunk = chunk.slice(0, bodyCloseIndex) + stateScript + chunk.slice(bodyCloseIndex);
+					}
+					injected = true;
 				}
-				injected = true;
 			}
-
-			rewriter.emitEndTag(tag);
-		});
-
-		const writeStream = new WritableStream({
-			write(chunk) {
-				return new Promise<void>((resolve) => {
-					rewriter.write(chunk, () => {
-						resolve();
-					});
-				});
-			},
-			close() {
-				rewriter.end();
-			},
-		});
-
-		const readStream = new ReadableStream({
-			pull(controller) {
-				if (rewriter.closed) {
-					controller.close();
-				}
-
-				const chunk = rewriter.read();
-				controller.enqueue(chunk);
-			},
-		});
-
-		const blob = await result.blob();
-
-		return blob.stream().pipeThrough({
-			writable: writeStream,
-			readable: readStream,
-		});
-	}
-
-	return new Response(await render(), result);
-});
+			controller.enqueue(chunk);
+		},
+	});
+}