[Markdoc] Support automatic image optimization with `experimental.assets` (#6630)

* wip: scrappy implementation. It works! 🥳

* chore: add code comments on inline utils

* fix: code cleanup, run on experimental.assets

* feat: support ~/assets alias

* fix: spoof `astro:assets` when outside experimental

* test: image paths in dev and prod

* feat: support any vite alias with ctx.resolve

* fix: avoid trying to process absolute paths

* fix: raise helpful error for invalid vite paths

* refactor: revert URL support on emitAsset

* chore: lint

* refactor: expose emitESMImage from assets base

* wip: why doesn't assets exist

* scary chore: make @astrojs/markdoc truly depend on astro

* fix: import emitESMImage straight from dist

* chore: remove type def from assets package

* chore: screw it, just ts ignore

* deps: rollup types

* refactor: optimize images during parse step

* chore: remove unneeded `.flat()`

* fix: use file-based relative paths

* fix: add back helpful error

* chore: changeset

* deps: move astro back to dev dep

* fix: put emit assets behind flag

* chore: change to markdoc patch

2 files changed, 155 insertions, 11 deletions

diff --git a/packages/integrations/markdoc/src/index.ts b/packages/integrations/markdoc/src/index.ts
index 70d005ee5..1d3556db7 100644
--- a/packages/integrations/markdoc/src/index.ts
+++ b/packages/integrations/markdoc/src/index.ts
@@ -1,9 +1,23 @@
-import type { Config } from '@markdoc/markdoc';
+import type {
+	Config as ReadonlyMarkdocConfig,
+	ConfigType as MarkdocConfig,
+	Node,
+} from '@markdoc/markdoc';
 import Markdoc from '@markdoc/markdoc';
 import type { AstroConfig, AstroIntegration, ContentEntryType, HookParameters } from 'astro';
 import fs from 'node:fs';
+import type * as rollup from 'rollup';
 import { fileURLToPath } from 'node:url';
-import { getAstroConfigPath, MarkdocError, parseFrontmatter } from './utils.js';
+import {
+	getAstroConfigPath,
+	isValidUrl,
+	MarkdocError,
+	parseFrontmatter,
+	prependForwardSlash,
+} from './utils.js';
+// @ts-expect-error Cannot find module 'astro/assets' or its corresponding type declarations.
+import { emitESMImage } from 'astro/assets';
+import type { Plugin as VitePlugin } from 'vite';
 
 type SetupHookParams = HookParameters<'astro:config:setup'> & {
 	// `contentEntryType` is not a public API
@@ -11,12 +25,24 @@ type SetupHookParams = HookParameters<'astro:config:setup'> & {
 	addContentEntryType: (contentEntryType: ContentEntryType) => void;
 };
 
-export default function markdoc(markdocConfig: Config = {}): AstroIntegration {
+export default function markdocIntegration(
+	userMarkdocConfig: ReadonlyMarkdocConfig = {}
+): AstroIntegration {
 	return {
 		name: '@astrojs/markdoc',
 		hooks: {
 			'astro:config:setup': async (params) => {
-				const { updateConfig, config, addContentEntryType } = params as SetupHookParams;
+				const {
+					updateConfig,
+					config: astroConfig,
+					addContentEntryType,
+				} = params as SetupHookParams;
+
+				updateConfig({
+					vite: {
+						plugins: [safeAssetsVirtualModulePlugin({ astroConfig })],
+					},
+				});
 
 				function getEntryInfo({ fileUrl, contents }: { fileUrl: URL; contents: string }) {
 					const parsed = parseFrontmatter(contents, fileURLToPath(fileUrl));
@@ -30,16 +56,44 @@ export default function markdoc(markdocConfig: Config = {}): AstroIntegration {
 				addContentEntryType({
 					extensions: ['.mdoc'],
 					getEntryInfo,
-					getRenderModule({ entry }) {
-						validateRenderProperties(markdocConfig, config);
+					async getRenderModule({ entry }) {
+						validateRenderProperties(userMarkdocConfig, astroConfig);
 						const ast = Markdoc.parse(entry.body);
-						const content = Markdoc.transform(ast, {
-							...markdocConfig,
+						const pluginContext = this;
+						const markdocConfig: MarkdocConfig = {
+							...userMarkdocConfig,
 							variables: {
-								...markdocConfig.variables,
+								...userMarkdocConfig.variables,
 								entry,
 							},
-						});
+						};
+
+						if (astroConfig.experimental?.assets) {
+							await emitOptimizedImages(ast.children, {
+								astroConfig,
+								pluginContext,
+								filePath: entry._internal.filePath,
+							});
+
+							markdocConfig.nodes ??= {};
+							markdocConfig.nodes.image = {
+								...Markdoc.nodes.image,
+								transform(node, config) {
+									const attributes = node.transformAttributes(config);
+									const children = node.transformChildren(config);
+
+									if (node.type === 'image' && '__optimizedSrc' in node.attributes) {
+										const { __optimizedSrc, ...rest } = node.attributes;
+										return new Markdoc.Tag('Image', { ...rest, src: __optimizedSrc }, children);
+									} else {
+										return new Markdoc.Tag('img', attributes, children);
+									}
+								},
+							};
+						}
+
+						const content = Markdoc.transform(ast, markdocConfig);
+
 						return {
 							code: `import { jsx as h } from 'astro/jsx-runtime';\nimport { Renderer } from '@astrojs/markdoc/components';\nconst transformedContent = ${JSON.stringify(
 								content
@@ -56,7 +110,54 @@ export default function markdoc(markdocConfig: Config = {}): AstroIntegration {
 	};
 }
 
-function validateRenderProperties(markdocConfig: Config, astroConfig: AstroConfig) {
+/**
+ * Emits optimized images, and appends the generated `src` to each AST node
+ * via the `__optimizedSrc` attribute.
+ */
+async function emitOptimizedImages(
+	nodeChildren: Node[],
+	ctx: {
+		pluginContext: rollup.PluginContext;
+		filePath: string;
+		astroConfig: AstroConfig;
+	}
+) {
+	for (const node of nodeChildren) {
+		if (
+			node.type === 'image' &&
+			typeof node.attributes.src === 'string' &&
+			shouldOptimizeImage(node.attributes.src)
+		) {
+			// Attempt to resolve source with Vite.
+			// This handles relative paths and configured aliases
+			const resolved = await ctx.pluginContext.resolve(node.attributes.src, ctx.filePath);
+
+			if (resolved?.id && fs.existsSync(new URL(prependForwardSlash(resolved.id), 'file://'))) {
+				const src = await emitESMImage(
+					resolved.id,
+					ctx.pluginContext.meta.watchMode,
+					ctx.pluginContext.emitFile,
+					{ config: ctx.astroConfig }
+				);
+				node.attributes.__optimizedSrc = src;
+			} else {
+				throw new MarkdocError({
+					message: `Could not resolve image ${JSON.stringify(
+						node.attributes.src
+					)} from ${JSON.stringify(ctx.filePath)}. Does the file exist?`,
+				});
+			}
+		}
+		await emitOptimizedImages(node.children, ctx);
+	}
+}
+
+function shouldOptimizeImage(src: string) {
+	// Optimize anything that is NOT external or an absolute path to `public/`
+	return !isValidUrl(src) && !src.startsWith('/');
+}
+
+function validateRenderProperties(markdocConfig: ReadonlyMarkdocConfig, astroConfig: AstroConfig) {
 	const tags = markdocConfig.tags ?? {};
 	const nodes = markdocConfig.nodes ?? {};
 
@@ -105,3 +206,37 @@ function validateRenderProperty({
 function isCapitalized(str: string) {
 	return str.length > 0 && str[0] === str[0].toUpperCase();
 }
+
+/**
+ * TODO: remove when `experimental.assets` is baselined.
+ *
+ * `astro:assets` will fail to resolve if the `experimental.assets` flag is not enabled.
+ * This ensures a fallback for the Markdoc renderer to safely import at the top level.
+ * @see ../components/TreeNode.ts
+ */
+function safeAssetsVirtualModulePlugin({
+	astroConfig,
+}: {
+	astroConfig: Pick<AstroConfig, 'experimental'>;
+}): VitePlugin {
+	const virtualModuleId = 'astro:markdoc-assets';
+	const resolvedVirtualModuleId = '\0' + virtualModuleId;
+
+	return {
+		name: 'astro:markdoc-safe-assets-virtual-module',
+		resolveId(id) {
+			if (id === virtualModuleId) {
+				return resolvedVirtualModuleId;
+			}
+		},
+		load(id) {
+			if (id !== resolvedVirtualModuleId) return;
+
+			if (astroConfig.experimental?.assets) {
+				return `export { Image } from 'astro:assets';`;
+			} else {
+				return `export const Image = () => { throw new Error('Cannot use the Image component without the \`experimental.assets\` flag.'); }`;
+			}
+		},
+	};
+}
diff --git a/packages/integrations/markdoc/src/utils.ts b/packages/integrations/markdoc/src/utils.ts
index 275c711f0..9d6e5af26 100644
--- a/packages/integrations/markdoc/src/utils.ts
+++ b/packages/integrations/markdoc/src/utils.ts
@@ -145,3 +145,12 @@ const componentsPropValidator = z.record(
 export function isCapitalized(str: string) {
 	return str.length > 0 && str[0] === str[0].toUpperCase();
 }
+
+export function isValidUrl(str: string): boolean {
+	try {
+		new URL(str);
+		return true;
+	} catch {
+		return false;
+	}
+}