From c9ae7b1b89e050900bbc111f29e8c5d95c26bf36 Mon Sep 17 00:00:00 2001
From: Bjorn Lu <bjornlu.dev@gmail.com>
Date: Thu, 12 Sep 2024 18:47:55 +0800
Subject: [PATCH] Handle MDX optimize for root hast node (#11975)

---
 .changeset/lovely-dolphins-draw.md            |  5 +++
 .../mdx/src/rehype-optimize-static.ts         | 10 ++++-
 .../fixtures/mdx-optimize/astro.config.mjs    | 40 ++++++++++++++++---
 .../mdx/test/mdx-optimize.test.js             | 11 +++++
 pnpm-lock.yaml                                |  2 -
 5 files changed, 58 insertions(+), 10 deletions(-)
 create mode 100644 .changeset/lovely-dolphins-draw.md

diff --git a/.changeset/lovely-dolphins-draw.md b/.changeset/lovely-dolphins-draw.md
new file mode 100644
index 000000000000..af4c2bd3c6dc
--- /dev/null
+++ b/.changeset/lovely-dolphins-draw.md
@@ -0,0 +1,5 @@
+---
+'@astrojs/mdx': patch
+---
+
+Handles nested root hast node when optimizing MDX
diff --git a/packages/integrations/mdx/src/rehype-optimize-static.ts b/packages/integrations/mdx/src/rehype-optimize-static.ts
index 80c28ab9ca91..eba31cae0d11 100644
--- a/packages/integrations/mdx/src/rehype-optimize-static.ts
+++ b/packages/integrations/mdx/src/rehype-optimize-static.ts
@@ -64,8 +64,14 @@ export const rehypeOptimizeStatic: RehypePlugin<[OptimizeOptions?]> = (options)
 		 * A non-static node causes all its parents to be non-optimizable
 		 */
 		const isNodeNonStatic = (node: Node) => {
-			// @ts-expect-error Access `.tagName` naively for perf
-			return node.type.startsWith('mdx') || ignoreElementNames.has(node.tagName);
+			return (
+				node.type.startsWith('mdx') ||
+				// @ts-expect-error `node` should never have `type: 'root'`, but in some cases plugins may inject it as children,
+				// which MDX will render as a fragment instead (an MDX fragment is a `mdxJsxFlowElement` type).
+				node.type === 'root' ||
+				// @ts-expect-error Access `.tagName` naively for perf
+				ignoreElementNames.has(node.tagName)
+			);
 		};
 
 		visit(tree as any, {
diff --git a/packages/integrations/mdx/test/fixtures/mdx-optimize/astro.config.mjs b/packages/integrations/mdx/test/fixtures/mdx-optimize/astro.config.mjs
index 204549479f5d..a3626c3a3129 100644
--- a/packages/integrations/mdx/test/fixtures/mdx-optimize/astro.config.mjs
+++ b/packages/integrations/mdx/test/fixtures/mdx-optimize/astro.config.mjs
@@ -1,9 +1,37 @@
 import mdx from '@astrojs/mdx';
 
 export default {
-	integrations: [mdx({
-		optimize: {
-			ignoreElementNames: ['strong']
-		}
-	})]
-}
+	integrations: [
+		mdx({
+			optimize: {
+				ignoreElementNames: ['strong'],
+			},
+		}),
+	],
+	markdown: {
+		rehypePlugins: [
+			() => {
+				return (tree) => {
+					tree.children.push({
+						type: 'root',
+						children: [
+							{
+								type: 'element',
+								tagName: 'p',
+								properties: {
+									id: 'injected-root-hast',
+								},
+								children: [
+									{
+										type: 'text',
+										value: 'Injected root hast from rehype plugin',
+									},
+								],
+							},
+						],
+					});
+				};
+			},
+		],
+	},
+};
diff --git a/packages/integrations/mdx/test/mdx-optimize.test.js b/packages/integrations/mdx/test/mdx-optimize.test.js
index 47acfdff80ce..c45ecda15015 100644
--- a/packages/integrations/mdx/test/mdx-optimize.test.js
+++ b/packages/integrations/mdx/test/mdx-optimize.test.js
@@ -47,4 +47,15 @@ describe('MDX optimize', () => {
 		assert.notEqual(blockquote, null);
 		assert.equal(blockquote.textContent.includes('I like pancakes'), true);
 	});
+
+	it('renders MDX with rehype plugin that incorrectly injects root hast node', async () => {
+		const html = await fixture.readFile('/import/index.html');
+		const { document } = parseHTML(html);
+
+		assert.doesNotMatch(html, /set:html=/);
+		assert.equal(
+			document.getElementById('injected-root-hast').textContent,
+			'Injected root hast from rehype plugin',
+		);
+	});
 });
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
index 9937787d9940..c658edeca2f2 100644
--- a/pnpm-lock.yaml
+++ b/pnpm-lock.yaml
@@ -8865,12 +8865,10 @@ packages:
 
   libsql@0.3.19:
     resolution: {integrity: sha512-Aj5cQ5uk/6fHdmeW0TiXK42FqUlwx7ytmMLPSaUQPin5HKKKuUPD62MAbN4OEweGBBI7q1BekoEN4gPUEL6MZA==}
-    cpu: [x64, arm64, wasm32]
     os: [darwin, linux, win32]
 
   libsql@0.4.1:
     resolution: {integrity: sha512-qZlR9Yu1zMBeLChzkE/cKfoKV3Esp9cn9Vx5Zirn4AVhDWPcjYhKwbtJcMuHehgk3mH+fJr9qW+3vesBWbQpBg==}
-    cpu: [x64, arm64, wasm32]
     os: [darwin, linux, win32]
 
   lilconfig@2.1.0: