feat(mdx): Add support for turning ![]() into <Image> (#6824)

12 files changed, 195 insertions, 7 deletions

diff --git a/.changeset/giant-squids-pull.md b/.changeset/giant-squids-pull.md
new file mode 100644
index 000000000..795bb6359
--- /dev/null
+++ b/.changeset/giant-squids-pull.md
@@ -0,0 +1,6 @@
+---
+'@astrojs/mdx': minor
+'@astrojs/markdown-remark': patch
+---
+
+Add support for using optimized and relative images in MDX files with `experimental.assets`
diff --git a/packages/integrations/mdx/src/plugins.ts b/packages/integrations/mdx/src/plugins.ts
index 12b8f2bd3..56fbbf837 100644
--- a/packages/integrations/mdx/src/plugins.ts
+++ b/packages/integrations/mdx/src/plugins.ts
@@ -1,4 +1,4 @@
-import { rehypeHeadingIds } from '@astrojs/markdown-remark';
+import { rehypeHeadingIds, remarkCollectImages } from '@astrojs/markdown-remark';
 import {
 	InvalidAstroDataError,
 	safelyGetAstroData,
@@ -16,6 +16,7 @@ import type { VFile } from 'vfile';
 import type { MdxOptions } from './index.js';
 import { rehypeInjectHeadingsExport } from './rehype-collect-headings.js';
 import rehypeMetaString from './rehype-meta-string.js';
+import { remarkImageToComponent } from './remark-images-to-component.js';
 import remarkPrism from './remark-prism.js';
 import remarkShiki from './remark-shiki.js';
 import { jsToTreeNode } from './utils.js';
@@ -99,7 +100,7 @@ export async function getRemarkPlugins(
 	mdxOptions: MdxOptions,
 	config: AstroConfig
 ): Promise<MdxRollupPluginOptions['remarkPlugins']> {
-	let remarkPlugins: PluggableList = [];
+	let remarkPlugins: PluggableList = [...(config.experimental.assets ? [remarkCollectImages, remarkImageToComponent] : [])];
 
 	if (!isPerformanceBenchmark) {
 		if (mdxOptions.gfm) {
diff --git a/packages/integrations/mdx/src/remark-images-to-component.ts b/packages/integrations/mdx/src/remark-images-to-component.ts
new file mode 100644
index 000000000..8a3166f49
--- /dev/null
+++ b/packages/integrations/mdx/src/remark-images-to-component.ts
@@ -0,0 +1,98 @@
+import type { MarkdownVFile } from '@astrojs/markdown-remark';
+import { type Image, type Parent } from 'mdast';
+import type { MdxJsxFlowElement, MdxjsEsm } from 'mdast-util-mdx';
+import { visit } from 'unist-util-visit';
+import { jsToTreeNode } from './utils.js';
+
+export function remarkImageToComponent() {
+	return function (tree: any, file: MarkdownVFile) {
+		if (!file.data.imagePaths) return;
+
+		const importsStatements: MdxjsEsm[] = [];
+		const importedImages = new Map<string, string>();
+
+		visit(tree, 'image', (node: Image, index: number | null, parent: Parent | null) => {
+			// Use the imagePaths set from the remark-collect-images so we don't have to duplicate the logic for
+			// checking if an image should be imported or not
+			if (file.data.imagePaths?.has(node.url)) {
+				let importName = importedImages.get(node.url);
+
+				// If we haven't already imported this image, add an import statement
+				if (!importName) {
+					importName = `__${importedImages.size}_${node.url.replace(/\W/g, '_')}__`;
+
+					importsStatements.push({
+						type: 'mdxjsEsm',
+						value: '',
+						data: {
+							estree: {
+								type: 'Program',
+								sourceType: 'module',
+								body: [
+									{
+										type: 'ImportDeclaration',
+										source: { type: 'Literal', value: node.url, raw: JSON.stringify(node.url) },
+										specifiers: [
+											{
+												type: 'ImportDefaultSpecifier',
+												local: { type: 'Identifier', name: importName },
+											},
+										],
+									},
+								],
+							},
+						},
+					});
+					importedImages.set(node.url, importName);
+				}
+
+				// Build a component that's equivalent to <Image src={importName} alt={node.alt} title={node.title} />
+				const componentElement: MdxJsxFlowElement = {
+					name: '__AstroImage__',
+					type: 'mdxJsxFlowElement',
+					attributes: [
+						{
+							name: 'src',
+							type: 'mdxJsxAttribute',
+							value: {
+								type: 'mdxJsxAttributeValueExpression',
+								value: importName,
+								data: {
+									estree: {
+										type: 'Program',
+										sourceType: 'module',
+										comments: [],
+										body: [
+											{
+												type: 'ExpressionStatement',
+												expression: { type: 'Identifier', name: importName },
+											},
+										],
+									},
+								},
+							},
+						},
+						{ name: 'alt', type: 'mdxJsxAttribute', value: node.alt || '' },
+					],
+					children: [],
+				};
+
+				if (node.title) {
+					componentElement.attributes.push({
+						type: 'mdxJsxAttribute',
+						name: 'title',
+						value: node.title,
+					});
+				}
+
+				parent!.children.splice(index!, 1, componentElement);
+			}
+		});
+
+		// Add all the import statements to the top of the file for the images
+		tree.children.unshift(...importsStatements);
+
+		// Add an import statement for the Astro Image component, we rename it to avoid conflicts
+		tree.children.unshift(jsToTreeNode(`import { Image as __AstroImage__ } from "astro:assets";`));
+	};
+}
diff --git a/packages/integrations/mdx/test/fixtures/mdx-images/astro.config.ts b/packages/integrations/mdx/test/fixtures/mdx-images/astro.config.ts
new file mode 100644
index 000000000..fe92bd37f
--- /dev/null
+++ b/packages/integrations/mdx/test/fixtures/mdx-images/astro.config.ts
@@ -0,0 +1,8 @@
+import mdx from '@astrojs/mdx';
+
+export default {
+	integrations: [mdx()],
+	experimental: {
+		assets: true
+	}
+}
diff --git a/packages/integrations/mdx/test/fixtures/mdx-images/package.json b/packages/integrations/mdx/test/fixtures/mdx-images/package.json
new file mode 100644
index 000000000..7ff215df1
--- /dev/null
+++ b/packages/integrations/mdx/test/fixtures/mdx-images/package.json
@@ -0,0 +1,9 @@
+{
+  "name": "@test/mdx-page",
+  "dependencies": {
+    "@astrojs/mdx": "workspace:*",
+    "astro": "workspace:*",
+    "react": "^18.2.0",
+    "react-dom": "^18.2.0"
+  }
+}
diff --git a/packages/integrations/mdx/test/fixtures/mdx-images/src/assets/houston in space.webp b/packages/integrations/mdx/test/fixtures/mdx-images/src/assets/houston in space.webp
new file mode 100644
index 000000000..3727bc508
--- /dev/null
+++ b/packages/integrations/mdx/test/fixtures/mdx-images/src/assets/houston in space.webp
diff --git a/packages/integrations/mdx/test/fixtures/mdx-images/src/assets/houston.webp b/packages/integrations/mdx/test/fixtures/mdx-images/src/assets/houston.webp
new file mode 100644
index 000000000..3727bc508
--- /dev/null
+++ b/packages/integrations/mdx/test/fixtures/mdx-images/src/assets/houston.webp
diff --git a/packages/integrations/mdx/test/fixtures/mdx-images/src/pages/index.mdx b/packages/integrations/mdx/test/fixtures/mdx-images/src/pages/index.mdx
new file mode 100644
index 000000000..b34d50b7c
--- /dev/null
+++ b/packages/integrations/mdx/test/fixtures/mdx-images/src/pages/index.mdx
@@ -0,0 +1,11 @@
+Image using a relative path:
+![Houston](../assets/houston.webp)
+
+Image using an aliased path:
+![Houston](~/assets/houston.webp)
+
+Image with a title:
+![Houston](~/assets/houston.webp "Houston title")
+
+Image with spaces in the path:
+![Houston](<~/assets/houston in space.webp>)
diff --git a/packages/integrations/mdx/test/mdx-images.test.js b/packages/integrations/mdx/test/mdx-images.test.js
new file mode 100644
index 000000000..c9c8e1f7c
--- /dev/null
+++ b/packages/integrations/mdx/test/mdx-images.test.js
@@ -0,0 +1,40 @@
+import { expect } from 'chai';
+import { parseHTML } from 'linkedom';
+import { loadFixture } from '../../../astro/test/test-utils.js';
+
+describe('MDX Page', () => {
+	let devServer;
+	let fixture;
+
+	before(async () => {
+		fixture = await loadFixture({
+			root: new URL('./fixtures/mdx-images/', import.meta.url),
+		});
+		devServer = await fixture.startDevServer();
+	});
+
+	after(async () => {
+		await devServer.stop();
+	});
+
+	describe('Optimized images in MDX', () => {
+		it('works', async () => {
+			const res = await fixture.fetch('/');
+			expect(res.status).to.equal(200);
+
+			const html = await res.text();
+			const { document } = parseHTML(html);
+
+			const imgs = document.getElementsByTagName('img');
+			expect(imgs.length).to.equal(4);
+			// Image using a relative path
+			expect(imgs.item(0).src.startsWith('/_image')).to.be.true;
+			// Image using an aliased path
+			expect(imgs.item(1).src.startsWith('/_image')).to.be.true;
+			// Image with title
+			expect(imgs.item(2).title).to.equal('Houston title');
+			// Image with spaces in the path
+			expect(imgs.item(3).src.startsWith('/_image')).to.be.true;
+		});
+	});
+});
diff --git a/packages/markdown/remark/src/index.ts b/packages/markdown/remark/src/index.ts
index f37b9ed68..0a21e1c98 100644
--- a/packages/markdown/remark/src/index.ts
+++ b/packages/markdown/remark/src/index.ts
@@ -8,7 +8,7 @@ import type {
 import { toRemarkInitializeAstroData } from './frontmatter-injection.js';
 import { loadPlugins } from './load-plugins.js';
 import { rehypeHeadingIds } from './rehype-collect-headings.js';
-import toRemarkCollectImages from './remark-collect-images.js';
+import { remarkCollectImages } from './remark-collect-images.js';
 import remarkPrism from './remark-prism.js';
 import scopedStyles from './remark-scoped-styles.js';
 import remarkShiki from './remark-shiki.js';
@@ -24,6 +24,7 @@ import { VFile } from 'vfile';
 import { rehypeImages } from './rehype-images.js';
 
 export { rehypeHeadingIds } from './rehype-collect-headings.js';
+export { remarkCollectImages } from './remark-collect-images.js';
 export * from './types.js';
 
 export const markdownConfigDefaults: Omit<Required<AstroMarkdownOptions>, 'drafts'> = {
@@ -96,7 +97,7 @@ export async function renderMarkdown(
 
 		if (opts.experimentalAssets) {
 			// Apply later in case user plugins resolve relative image paths
-			parser.use([toRemarkCollectImages()]);
+			parser.use([remarkCollectImages]);
 		}
 	}
 
diff --git a/packages/markdown/remark/src/remark-collect-images.ts b/packages/markdown/remark/src/remark-collect-images.ts
index 470b770ed..0f1eb59f7 100644
--- a/packages/markdown/remark/src/remark-collect-images.ts
+++ b/packages/markdown/remark/src/remark-collect-images.ts
@@ -2,9 +2,8 @@ import type { Image } from 'mdast';
 import { visit } from 'unist-util-visit';
 import type { MarkdownVFile } from './types';
 
-export default function toRemarkCollectImages() {
-	return () =>
-		async function (tree: any, vfile: MarkdownVFile) {
+export function remarkCollectImages() {
+	return function (tree: any, vfile: MarkdownVFile) {
 			if (typeof vfile?.path !== 'string') return;
 
 			const imagePaths = new Set<string>();
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
index c0fa7f3d8..d6f7ec104 100644
--- a/pnpm-lock.yaml
+++ b/pnpm-lock.yaml
@@ -4085,6 +4085,21 @@ importers:
         specifier: ^4.1.0
         version: 4.1.2
 
+  packages/integrations/mdx/test/fixtures/mdx-images:
+    dependencies:
+      '@astrojs/mdx':
+        specifier: workspace:*
+        version: link:../../..
+      astro:
+        specifier: workspace:*
+        version: link:../../../../../astro
+      react:
+        specifier: ^18.2.0
+        version: 18.2.0
+      react-dom:
+        specifier: ^18.2.0
+        version: 18.2.0(react@18.2.0)
+
   packages/integrations/mdx/test/fixtures/mdx-infinite-loop:
     dependencies:
       '@astrojs/mdx':