Handle metadata for MDX files during build (#3915)

* fix: metadata handling for MDX files

* chore: add changeset

* chore: update mdx example

* fix: protect against infinite loops in jsx-runtime, properly hook console.error

* chore: remove unused import

* feat(mdx): support `client:only`

* fix: prefer Symbol.for

* fix(jsx): handle vnode check properly

* chore: appease ts

Co-authored-by: Nate Moore <nate@astro.build>

.changeset/purple-flowers-arrive.md

@@ -0,0 +1,5 @@
+---
+'astro': patch
+---
+
+Fix metadata handling for building MDX files

examples/with-mdx/src/pages/index.mdx

@@ -14,4 +14,4 @@ Written by: {new Intl.ListFormat('en').format(authors.map(d => d.name))}.
 
 Published on: {new Intl.DateTimeFormat('en', {dateStyle: 'long'}).format(published)}.
 
-<Counter>## This is a counter!</Counter>
+<Counter client:idle>This is a **counter**!</Counter>

packages/astro/src/jsx-runtime/index.ts

@@ -3,9 +3,9 @@ import { Fragment, markHTMLString } from '../runtime/server/index.js';
 const AstroJSX = 'astro:jsx';
 const Empty = Symbol('empty');
 
-interface AstroVNode {
+export interface AstroVNode {
 	[AstroJSX]: boolean;
-	type: string | ((...args: any) => any) | typeof Fragment;
+	type: string | ((...args: any) => any);
 	props: Record<string, any>;
 }
 

packages/astro/src/jsx/babel.ts

@@ -1,5 +1,8 @@
-import type { PluginObj } from '@babel/core';
+import type { PluginMetadata } from '../vite-plugin-astro/types';
+import type { PluginObj, NodePath } from '@babel/core';
 import * as t from '@babel/types';
+import { pathToFileURL } from 'node:url'
+import { ClientOnlyPlaceholder } from '../runtime/server/index.js';
 
 function isComponent(tagName: string) {
 	return (
@@ -18,6 +21,15 @@ function hasClientDirective(node: t.JSXElement) {
 	return false;
 }
 
+function isClientOnlyComponent(node: t.JSXElement) {
+	for (const attr of node.openingElement.attributes) {
+		if (attr.type === 'JSXAttribute' && attr.name.type === 'JSXNamespacedName') {
+			return jsxAttributeToString(attr) === 'client:only'; 
+		}
+	}
+	return false;
+}
+
 function getTagName(tag: t.JSXElement) {
 	const jsxName = tag.openingElement.name;
 	return jsxElementNameToString(jsxName);
@@ -40,28 +52,55 @@ function jsxAttributeToString(attr: t.JSXAttribute): string {
 	return `${attr.name.name}`;
 }
 
-function addClientMetadata(node: t.JSXElement, meta: { path: string; name: string }) {
+function addClientMetadata(node: t.JSXElement, meta: { resolvedPath: string; path: string; name: string }) {
 	const existingAttributes = node.openingElement.attributes.map((attr) =>
 		t.isJSXAttribute(attr) ? jsxAttributeToString(attr) : null
 	);
 	if (!existingAttributes.find((attr) => attr === 'client:component-path')) {
 		const componentPath = t.jsxAttribute(
 			t.jsxNamespacedName(t.jsxIdentifier('client'), t.jsxIdentifier('component-path')),
-			!meta.path.startsWith('.')
-				? t.stringLiteral(meta.path)
-				: t.jsxExpressionContainer(
-						t.binaryExpression(
-							'+',
-							t.stringLiteral('/@fs'),
-							t.memberExpression(
-								t.newExpression(t.identifier('URL'), [
-									t.stringLiteral(meta.path),
-									t.identifier('import.meta.url'),
-								]),
-								t.identifier('pathname')
-							)
-						)
-				  )
+			t.stringLiteral(meta.resolvedPath)
+		);
+		node.openingElement.attributes.push(componentPath);
+	}
+	if (!existingAttributes.find((attr) => attr === 'client:component-export')) {
+		if (meta.name === '*') {
+			meta.name = getTagName(node).split('.').at(1)!;
+		}
+		const componentExport = t.jsxAttribute(
+			t.jsxNamespacedName(t.jsxIdentifier('client'), t.jsxIdentifier('component-export')),
+			t.stringLiteral(meta.name)
+		);
+		node.openingElement.attributes.push(componentExport);
+	}
+	if (!existingAttributes.find((attr) => attr === 'client:component-hydration')) {
+		const staticMarker = t.jsxAttribute(
+			t.jsxNamespacedName(t.jsxIdentifier('client'), t.jsxIdentifier('component-hydration'))
+		);
+		node.openingElement.attributes.push(staticMarker);
+	}
+}
+
+function addClientOnlyMetadata(node: t.JSXElement, meta: { resolvedPath: string; path: string; name: string }) {
+	const tagName = getTagName(node);
+	node.openingElement = t.jsxOpeningElement(t.jsxIdentifier(ClientOnlyPlaceholder), node.openingElement.attributes)
+	if (node.closingElement) {
+		node.closingElement = t.jsxClosingElement(t.jsxIdentifier(ClientOnlyPlaceholder))
+	}
+	const existingAttributes = node.openingElement.attributes.map((attr) =>
+		t.isJSXAttribute(attr) ? jsxAttributeToString(attr) : null
+	);
+	if (!existingAttributes.find((attr) => attr === 'client:display-name')) {
+		const displayName = t.jsxAttribute(
+			t.jsxNamespacedName(t.jsxIdentifier('client'), t.jsxIdentifier('display-name')),
+			t.stringLiteral(tagName)
+		);
+		node.openingElement.attributes.push(displayName);
+	}
+	if (!existingAttributes.find((attr) => attr === 'client:component-path')) {
+		const componentPath = t.jsxAttribute(
+			t.jsxNamespacedName(t.jsxIdentifier('client'), t.jsxIdentifier('component-path')),
+			t.stringLiteral(meta.resolvedPath)
 		);
 		node.openingElement.attributes.push(componentPath);
 	}
@@ -86,7 +125,15 @@ function addClientMetadata(node: t.JSXElement, meta: { path: string; name: strin
 export default function astroJSX(): PluginObj {
 	return {
 		visitor: {
-			Program(path) {
+			Program: {
+				enter(path, state) {
+					if (!(state.file.metadata as PluginMetadata).astro) {
+						(state.file.metadata as PluginMetadata).astro = {
+							clientOnlyComponents: [],
+							hydratedComponents: [],
+							scripts: [],
+						}
+					}
 					path.node.body.splice(
 						0,
 						0,
@@ -95,6 +142,7 @@ export default function astroJSX(): PluginObj {
 							t.stringLiteral('astro/jsx-runtime')
 						)
 					);
+				}
 			},
 			ImportDeclaration(path, state) {
 				const source = path.node.source.value;
@@ -127,9 +175,11 @@ export default function astroJSX(): PluginObj {
 				const tagName = getTagName(parentNode);
 				if (!isComponent(tagName)) return;
 				if (!hasClientDirective(parentNode)) return;
+				const isClientOnly = isClientOnlyComponent(parentNode);
+				if (tagName === ClientOnlyPlaceholder) return;
 
 				const imports = state.get('imports') ?? new Map();
-				const namespace = getTagName(parentNode).split('.');
+				const namespace = tagName.split('.');
 				for (const [source, specs] of imports) {
 					for (const { imported, local } of specs) {
 						const reference = path.referencesImport(source, imported);
@@ -143,10 +193,38 @@ export default function astroJSX(): PluginObj {
 						}
 					}
 				}
-				// TODO: map unmatched identifiers back to imports if possible
+
 				const meta = path.getData('import');
 				if (meta) {
+					let resolvedPath: string;
+					if (meta.path.startsWith('.')) {
+						const fileURL = pathToFileURL(state.filename!);
+						resolvedPath = `/@fs${new URL(meta.path, fileURL).pathname}`;
+						if (resolvedPath.endsWith('.jsx')) {
+							resolvedPath = resolvedPath.slice(0, -4);
+						}
+					} else {
+						resolvedPath = meta.path;
+					}
+					if (isClientOnly) {
+						(state.file.metadata as PluginMetadata).astro.clientOnlyComponents.push({
+							exportName: meta.name,
+							specifier: meta.name,
+							resolvedPath
+						})
+
+						meta.resolvedPath = resolvedPath;
+						addClientOnlyMetadata(parentNode, meta);
+					} else {
+						(state.file.metadata as PluginMetadata).astro.hydratedComponents.push({
+							exportName: meta.name,
+							specifier: meta.name,
+							resolvedPath
+						})
+
+						meta.resolvedPath = resolvedPath;
 						addClientMetadata(parentNode, meta);
+					}
 				} else {
 					throw new Error(
 						`Unable to match <${getTagName(

packages/astro/src/jsx/server.ts

@@ -34,10 +34,8 @@ export async function renderToStaticMarkup(
 	}
 
 	const { result } = this;
-	try {
 	const html = await renderJSX(result, jsx(Component, { ...props, ...slots, children }));
 	return { html };
-	} catch (e) {}
 }
 
 export default {

packages/astro/src/runtime/server/hydration.ts

@@ -58,6 +58,9 @@ export function extractDirectives(inputProps: Record<string | number, any>): Ext
 				case 'client:component-hydration': {
 					break;
 				}
+				case 'client:display-name': {
+					break;
+				}
 				default: {
 					extracted.hydration.directive = key.split(':')[1];
 					extracted.hydration.value = value;

packages/astro/src/runtime/server/index.ts

@@ -155,7 +155,8 @@ export function mergeSlots(...slotted: unknown[]) {
 	return slots;
 }
 
-export const Fragment = Symbol('Astro.Fragment');
+export const Fragment = Symbol.for('astro:fragment');
+export const ClientOnlyPlaceholder = 'astro-client-only';
 
 function guessRenderers(componentUrl?: string): string[] {
 	const extname = componentUrl?.split('.').pop();

packages/astro/src/runtime/server/jsx.ts

@@ -1,7 +1,9 @@
+/* eslint-disable no-console */
+import { SSRResult } from '../../@types/astro.js';
 import { AstroJSX, isVNode } from '../../jsx-runtime/index.js';
 import {
 	escapeHTML,
-	Fragment,
+	ClientOnlyPlaceholder,
 	HTMLString,
 	markHTMLString,
 	renderComponent,
@@ -10,7 +12,11 @@ import {
 	voidElementNames,
 } from './index.js';
 
-export async function renderJSX(result: any, vnode: any): Promise<any> {
+const skipAstroJSXCheck = new WeakSet();
+let originalConsoleError: any;
+let consoleFilterRefs = 0;
+
+export async function renderJSX(result: SSRResult, vnode: any): Promise<any> {
 	switch (true) {
 		case vnode instanceof HTMLString:
 			return vnode;
@@ -18,13 +24,16 @@ export async function renderJSX(result: any, vnode: any): Promise<any> {
 			return markHTMLString(escapeHTML(vnode));
 		case !vnode && vnode !== 0:
 			return '';
-		case vnode.type === Fragment:
-			return renderJSX(result, vnode.props.children);
 		case Array.isArray(vnode):
 			return markHTMLString(
 				(await Promise.all(vnode.map((v: any) => renderJSX(result, v)))).join('')
 			);
-		case vnode.type.isAstroComponentFactory: {
+	}
+	if (isVNode(vnode)) {
+		switch (true) {
+			case (vnode.type as any) === Symbol.for('astro:fragment'):
+				return renderJSX(result, vnode.props.children);
+			case (vnode.type as any).isAstroComponentFactory: {
 				let props: Record<string, any> = {};
 				let slots: Record<string, any> = {};
 				for (const [key, value] of Object.entries(vnode.props ?? {})) {
@@ -37,24 +46,34 @@ export async function renderJSX(result: any, vnode: any): Promise<any> {
 						props[key] = value;
 					}
 				}
-			return await renderToString(result, vnode.type, props, slots);
+				return await renderToString(result, vnode.type as any, props, slots);
 			}
+			case !vnode.type && (vnode.type as any) !== 0:
+				return '';
+			case typeof vnode.type === 'string' && vnode.type !== ClientOnlyPlaceholder:
+				return await renderElement(result, vnode.type as string, vnode.props ?? {});
 		}
-	if (vnode[AstroJSX]) {
-		if (!vnode.type && vnode.type !== 0) return '';
-		if (typeof vnode.type === 'string') {
-			return await renderElement(result, vnode.type, vnode.props ?? {});
+
+		if (vnode.type) {
+			if (typeof vnode.type === 'function' && vnode.props['server:root']) {
+				const output = await vnode.type(vnode.props ?? {});
+				return await renderJSX(result, output);
 			}
-		if (!!vnode.type) {
+			if (typeof vnode.type === 'function' && !skipAstroJSXCheck.has(vnode.type)) {
+				useConsoleFilter();
 				try {
-				// TODO: silence Invalid hook call warning from React
 					const output = await vnode.type(vnode.props ?? {});
 					if (output && output[AstroJSX]) {
 						return await renderJSX(result, output);
 					} else if (!output) {
 						return await renderJSX(result, output);
 					}
-			} catch (e) {}
+				} catch (e) {
+					skipAstroJSXCheck.add(vnode.type);
+				} finally {
+					finishUsingConsoleFilter();
+				}
+			}
 
 			const { children = null, ...props } = vnode.props ?? {};
 			const slots: Record<string, any> = {
@@ -78,9 +97,20 @@ export async function renderJSX(result: any, vnode: any): Promise<any> {
 			for (const [key, value] of Object.entries(slots)) {
 				slots[key] = () => renderJSX(result, value);
 			}
-			return markHTMLString(
-				await renderComponent(result, vnode.type.name, vnode.type, props, slots)
+
+			let output: string | AsyncIterable<string>;
+			if (vnode.type === ClientOnlyPlaceholder && vnode.props['client:only']) {
+				output = await renderComponent(
+					result,
+					vnode.props['client:display-name'] ?? '',
+					null,
+					props,
+					slots
 				);
+			} else {
+				output = await renderComponent(result, typeof vnode.type === 'function' ? vnode.type.name : vnode.type, vnode.type, props, slots);
+			}
+			return markHTMLString(output);
 		}
 	}
 	// numbers, plain objects, etc
@@ -100,3 +130,57 @@ async function renderElement(
 		)}`
 	);
 }
+
+/**
+ * Reduces console noise by filtering known non-problematic errors.
+ *
+ * Performs reference counting to allow parallel usage from async code.
+ *
+ * To stop filtering, please ensure that there always is a matching call
+ * to `finishUsingConsoleFilter` afterwards.
+ */
+function useConsoleFilter() {
+	consoleFilterRefs++;
+
+	if (!originalConsoleError) {
+		// eslint-disable-next-line no-console
+		originalConsoleError = console.error;
+
+		try {
+			// eslint-disable-next-line no-console
+			console.error = filteredConsoleError;
+		} catch (error) {
+			// If we're unable to hook `console.error`, just accept it
+		}
+	}
+}
+
+/**
+ * Indicates that the filter installed by `useConsoleFilter`
+ * is no longer needed by the calling code.
+ */
+function finishUsingConsoleFilter() {
+	consoleFilterRefs--;
+
+	// Note: Instead of reverting `console.error` back to the original
+	// when the reference counter reaches 0, we leave our hook installed
+	// to prevent potential race conditions once `check` is made async
+}
+
+/**
+ * Hook/wrapper function for the global `console.error` function.
+ *
+ * Ignores known non-problematic errors while any code is using the console filter.
+ * Otherwise, simply forwards all arguments to the original function.
+ */
+function filteredConsoleError(msg: any, ...rest: any[]) {
+	if (consoleFilterRefs > 0 && typeof msg === 'string') {
+		// In `check`, we attempt to render JSX components through Preact.
+		// When attempting this on a React component, React may output
+		// the following error, which we can safely filter out:
+		const isKnownReactHookError =
+			msg.includes('Warning: Invalid hook call.') &&
+			msg.includes('https://reactjs.org/link/invalid-hook-call');
+		if (isKnownReactHookError) return;
+	}
+}

packages/astro/src/vite-plugin-jsx/index.ts

@@ -2,6 +2,7 @@ import type { TransformResult } from 'rollup';
 import type { Plugin, ResolvedConfig } from 'vite';
 import type { AstroConfig, AstroRenderer } from '../@types/astro';
 import type { LogOptions } from '../core/logger/core.js';
+import type { PluginMetadata } from '../vite-plugin-astro/types';
 
 import babel from '@babel/core';
 import * as eslexer from 'es-module-lexer';
@@ -70,6 +71,23 @@ async function transformJSX({
 	// TODO: Be more strict about bad return values here.
 	// Should we throw an error instead? Should we never return `{code: ""}`?
 	if (!result) return null;
+
+	if (renderer.name === 'astro:jsx') {
+		const { astro } = result.metadata as unknown as PluginMetadata;
+		return {
+			code: result.code || '',
+			map: result.map,
+			meta: {
+				astro,
+				vite: {
+					// Setting this vite metadata to `ts` causes Vite to resolve .js
+					// extensions to .ts files.
+					lang: 'ts',
+				},
+			},
+		};
+	}
+	
 	return {
 		code: result.code || '',
 		map: result.map,