Markdoc - Shiki (#7187)

* chore: remove unused util

* chore: changeset

* deps: shiki

* wip: first stab at shiki markdoc config

* feat: get shiki working!

* refactor: return HTML string directly from transform

* chore: move shiki to markdoc dev dep

* refactor: use async cache with clear docs on why

* test: transform units with Shiki config options

* refactor: switch to `extends` model

* refactor: nodes/ -> extensions/

* feat: raise friendly error for Promise extensions

* docs: README

* chore: lint

* chore: dead file

* chore: lowercase for fuzzy find please

* fix: bad ctx spread

* chore: clean up cache, add shiki imp error

* chore: add shiki to optional peer deps

* chore: hoist those consts

* docs: more explicit "install shiki now please"

Co-authored-by: Sarah Rainsberger <sarah@rainsberger.ca>

* oops bad find and replace

* chore: update changeset

* nit: period haunts me

---------

Co-authored-by: Sarah Rainsberger <sarah@rainsberger.ca>

13 files changed, 383 insertions, 35 deletions

diff --git a/.changeset/eleven-tables-speak.md b/.changeset/eleven-tables-speak.md
new file mode 100644
index 000000000..6ff1474c7
--- /dev/null
+++ b/.changeset/eleven-tables-speak.md
@@ -0,0 +1,17 @@
+---
+'@astrojs/markdoc': patch
+---
+
+Add support for syntax highlighting with Shiki. Install `shiki` in your project with `npm i shiki`, and apply to your Markdoc config using the `extends` option:
+
+```js
+// markdoc.config.mjs
+import { defineMarkdocConfig, shiki } from '@astrojs/markdoc/config';
+export default defineMarkdocConfig({
+  extends: [
+    await shiki({ /** Shiki config options */ }),
+  ],
+})
+```
+
+Learn more in the [`@astrojs/markdoc` README.](https://docs.astro.build/en/guides/integrations-guide/markdoc/#syntax-highlighting)
diff --git a/packages/astro/src/runtime/server/index.ts b/packages/astro/src/runtime/server/index.ts
index 1f1e1e97b..021e55a56 100644
--- a/packages/astro/src/runtime/server/index.ts
+++ b/packages/astro/src/runtime/server/index.ts
@@ -1,7 +1,14 @@
 export { createComponent } from './astro-component.js';
 export { createAstro } from './astro-global.js';
 export { renderEndpoint } from './endpoint.js';
-export { escapeHTML, HTMLBytes, HTMLString, markHTMLString, unescapeHTML } from './escape.js';
+export {
+	escapeHTML,
+	HTMLBytes,
+	HTMLString,
+	markHTMLString,
+	unescapeHTML,
+	isHTMLString,
+} from './escape.js';
 export { renderJSX } from './jsx.js';
 export {
 	addAttribute,
diff --git a/packages/integrations/markdoc/README.md b/packages/integrations/markdoc/README.md
index e3cec5499..815f0420b 100644
--- a/packages/integrations/markdoc/README.md
+++ b/packages/integrations/markdoc/README.md
@@ -203,6 +203,41 @@ export default defineMarkdocConfig({
 })
 ```
 
+### Syntax highlighting
+
+`@astrojs/markdoc` provides a [Shiki](https://github.com/shikijs/shiki) extension to highlight your code blocks.
+
+To use this extension, you must separately install `shiki` as a dependency:
+
+```bash
+npm i shiki
+```
+
+Then, apply the `shiki()` extension to your Markdoc config using the `extends` property. You can optionally pass a shiki configuration object:
+
+```js
+// markdoc.config.mjs
+import { defineMarkdocConfig, shiki } from '@astrojs/markdoc/config';
+
+export default defineMarkdocConfig({
+  extends: [
+    await shiki({
+      // Choose from Shiki's built-in themes (or add your own)
+      // Default: 'github-dark'
+      // https://github.com/shikijs/shiki/blob/main/docs/themes.md
+      theme: 'dracula',
+      // Enable word wrap to prevent horizontal scrolling
+      // Default: false
+      wrap: true,
+      // Pass custom languages
+      // Note: Shiki has countless langs built-in, including `.astro`!
+      // https://github.com/shikijs/shiki/blob/main/docs/languages.md
+      langs: [],
+    })
+  ],
+})
+```
+
 ### Access frontmatter and content collection information from your templates
 
 You can access content collection information from your Markdoc templates using the `$entry` variable. This includes the entry `slug`, `collection` name, and frontmatter `data` parsed by your content collection schema (if any). This example renders the `title` frontmatter property as a heading:
diff --git a/packages/integrations/markdoc/components/TreeNode.ts b/packages/integrations/markdoc/components/TreeNode.ts
index a60597a0d..d12180a18 100644
--- a/packages/integrations/markdoc/components/TreeNode.ts
+++ b/packages/integrations/markdoc/components/TreeNode.ts
@@ -2,12 +2,18 @@ import type { AstroInstance } from 'astro';
 import { Fragment } from 'astro/jsx-runtime';
 import type { RenderableTreeNode } from '@markdoc/markdoc';
 import Markdoc from '@markdoc/markdoc';
-import { createComponent, renderComponent, render } from 'astro/runtime/server/index.js';
+import {
+	createComponent,
+	renderComponent,
+	render,
+	HTMLString,
+	isHTMLString,
+} from 'astro/runtime/server/index.js';
 
 export type TreeNode =
 	| {
 			type: 'text';
-			content: string;
+			content: string | HTMLString;
 	  }
 	| {
 			type: 'component';
@@ -25,6 +31,7 @@ export type TreeNode =
 export const ComponentNode = createComponent({
 	factory(result: any, { treeNode }: { treeNode: TreeNode }) {
 		if (treeNode.type === 'text') return render`${treeNode.content}`;
+
 		const slots = {
 			default: () =>
 				render`${treeNode.children.map((child) =>
@@ -46,7 +53,9 @@ export const ComponentNode = createComponent({
 });
 
 export function createTreeNode(node: RenderableTreeNode | RenderableTreeNode[]): TreeNode {
-	if (typeof node === 'string' || typeof node === 'number') {
+	if (isHTMLString(node)) {
+		return { type: 'text', content: node as HTMLString };
+	} else if (typeof node === 'string' || typeof node === 'number') {
 		return { type: 'text', content: String(node) };
 	} else if (Array.isArray(node)) {
 		return {
diff --git a/packages/integrations/markdoc/package.json b/packages/integrations/markdoc/package.json
index f031c8f6c..2086073ad 100644
--- a/packages/integrations/markdoc/package.json
+++ b/packages/integrations/markdoc/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@astrojs/markdoc",
-  "description": "Add support for Markdoc pages in your Astro site",
+  "description": "Add support for Markdoc in your Astro site",
   "version": "0.2.3",
   "type": "module",
   "types": "./dist/index.d.ts",
@@ -47,7 +47,13 @@
     "zod": "^3.17.3"
   },
   "peerDependencies": {
-    "astro": "workspace:^2.5.5"
+    "astro": "workspace:^2.5.5",
+    "shiki": "^0.14.1"
+  },
+  "peerDependenciesMeta": {
+    "shiki": {
+      "optional": true
+    }
   },
   "devDependencies": {
     "@astrojs/markdown-remark": "^2.2.1",
@@ -61,6 +67,7 @@
     "linkedom": "^0.14.12",
     "mocha": "^9.2.2",
     "rollup": "^3.20.1",
+    "shiki": "^0.14.1",
     "vite": "^4.3.1"
   },
   "engines": {
diff --git a/packages/integrations/markdoc/src/config.ts b/packages/integrations/markdoc/src/config.ts
index f8943ba1a..a8f202424 100644
--- a/packages/integrations/markdoc/src/config.ts
+++ b/packages/integrations/markdoc/src/config.ts
@@ -1,10 +1,19 @@
 import type { ConfigType as MarkdocConfig } from '@markdoc/markdoc';
 import _Markdoc from '@markdoc/markdoc';
-import { nodes as astroNodes } from './nodes/index.js';
+import { heading } from './heading-ids.js';
+
+export type AstroMarkdocConfig<C extends Record<string, any> = Record<string, any>> =
+	MarkdocConfig & {
+		ctx?: C;
+		extends?: ResolvedAstroMarkdocConfig[];
+	};
+
+export type ResolvedAstroMarkdocConfig = Omit<AstroMarkdocConfig, 'extends'>;
 
 export const Markdoc = _Markdoc;
-export const nodes = { ...Markdoc.nodes, ...astroNodes };
+export const nodes = { ...Markdoc.nodes, heading };
+export { shiki } from './extensions/shiki.js';
 
-export function defineMarkdocConfig(config: MarkdocConfig): MarkdocConfig {
+export function defineMarkdocConfig(config: AstroMarkdocConfig): AstroMarkdocConfig {
 	return config;
 }
diff --git a/packages/integrations/markdoc/src/extensions/shiki.ts b/packages/integrations/markdoc/src/extensions/shiki.ts
new file mode 100644
index 000000000..96d91d541
--- /dev/null
+++ b/packages/integrations/markdoc/src/extensions/shiki.ts
@@ -0,0 +1,138 @@
+// @ts-expect-error Cannot find module 'astro/runtime/server/index.js' or its corresponding type declarations.
+import { unescapeHTML } from 'astro/runtime/server/index.js';
+import type { ShikiConfig } from 'astro';
+import type * as shikiTypes from 'shiki';
+import type { AstroMarkdocConfig } from '../config.js';
+import Markdoc from '@markdoc/markdoc';
+import { MarkdocError } from '../utils.js';
+
+// Map of old theme names to new names to preserve compatibility when we upgrade shiki
+const compatThemes: Record<string, string> = {
+	'material-darker': 'material-theme-darker',
+	'material-default': 'material-theme',
+	'material-lighter': 'material-theme-lighter',
+	'material-ocean': 'material-theme-ocean',
+	'material-palenight': 'material-theme-palenight',
+};
+
+const normalizeTheme = (theme: string | shikiTypes.IShikiTheme) => {
+	if (typeof theme === 'string') {
+		return compatThemes[theme] || theme;
+	} else if (compatThemes[theme.name]) {
+		return { ...theme, name: compatThemes[theme.name] };
+	} else {
+		return theme;
+	}
+};
+
+const ASTRO_COLOR_REPLACEMENTS = {
+	'#000001': 'var(--astro-code-color-text)',
+	'#000002': 'var(--astro-code-color-background)',
+	'#000004': 'var(--astro-code-token-constant)',
+	'#000005': 'var(--astro-code-token-string)',
+	'#000006': 'var(--astro-code-token-comment)',
+	'#000007': 'var(--astro-code-token-keyword)',
+	'#000008': 'var(--astro-code-token-parameter)',
+	'#000009': 'var(--astro-code-token-function)',
+	'#000010': 'var(--astro-code-token-string-expression)',
+	'#000011': 'var(--astro-code-token-punctuation)',
+	'#000012': 'var(--astro-code-token-link)',
+};
+
+const PRE_SELECTOR = /<pre class="(.*?)shiki(.*?)"/;
+const LINE_SELECTOR = /<span class="line"><span style="(.*?)">([\+|\-])/g;
+const INLINE_STYLE_SELECTOR = /style="(.*?)"/;
+
+/**
+ * Note: cache only needed for dev server reloads, internal test suites, and manual calls to `Markdoc.transform` by the user.
+ * Otherwise, `shiki()` is only called once per build, NOT once per page, so a cache isn't needed!
+ */
+const highlighterCache = new Map<string, shikiTypes.Highlighter>();
+
+export async function shiki({
+	langs = [],
+	theme = 'github-dark',
+	wrap = false,
+}: ShikiConfig = {}): Promise<AstroMarkdocConfig> {
+	let getHighlighter: (options: shikiTypes.HighlighterOptions) => Promise<shikiTypes.Highlighter>;
+	try {
+		getHighlighter = (await import('shiki')).getHighlighter;
+	} catch {
+		throw new MarkdocError({
+			message: 'Shiki is not installed. Run `npm install shiki` to use the `shiki` extension.',
+		});
+	}
+	theme = normalizeTheme(theme);
+
+	const cacheID: string = typeof theme === 'string' ? theme : theme.name;
+	if (!highlighterCache.has(cacheID)) {
+		highlighterCache.set(
+			cacheID,
+			await getHighlighter({ theme }).then((hl) => {
+				hl.setColorReplacements(ASTRO_COLOR_REPLACEMENTS);
+				return hl;
+			})
+		);
+	}
+	const highlighter = highlighterCache.get(cacheID)!;
+
+	for (const lang of langs) {
+		await highlighter.loadLanguage(lang);
+	}
+	return {
+		nodes: {
+			fence: {
+				attributes: Markdoc.nodes.fence.attributes!,
+				transform({ attributes }) {
+					let lang: string;
+
+					if (typeof attributes.language === 'string') {
+						const langExists = highlighter
+							.getLoadedLanguages()
+							.includes(attributes.language as any);
+						if (langExists) {
+							lang = attributes.language;
+						} else {
+							// eslint-disable-next-line no-console
+							console.warn(
+								`[Shiki highlighter] The language "${attributes.language}" doesn't exist, falling back to plaintext.`
+							);
+							lang = 'plaintext';
+						}
+					} else {
+						lang = 'plaintext';
+					}
+
+					let html = highlighter.codeToHtml(attributes.content, { lang });
+
+					// Q: Could these regexes match on a user's inputted code blocks?
+					// A: Nope! All rendered HTML is properly escaped.
+					// Ex. If a user typed `<span class="line"` into a code block,
+					// It would become this before hitting our regexes:
+					// &lt;span class=&quot;line&quot;
+
+					html = html.replace(PRE_SELECTOR, `<pre class="$1astro-code$2"`);
+					// Add "user-select: none;" for "+"/"-" diff symbols
+					if (attributes.language === 'diff') {
+						html = html.replace(
+							LINE_SELECTOR,
+							'<span class="line"><span style="$1"><span style="user-select: none;">$2</span>'
+						);
+					}
+
+					if (wrap === false) {
+						html = html.replace(INLINE_STYLE_SELECTOR, 'style="$1; overflow-x: auto;"');
+					} else if (wrap === true) {
+						html = html.replace(
+							INLINE_STYLE_SELECTOR,
+							'style="$1; overflow-x: auto; white-space: pre-wrap; word-wrap: break-word;"'
+						);
+					}
+
+					// Use `unescapeHTML` to return `HTMLString` for Astro renderer to inline as HTML
+					return unescapeHTML(html);
+				},
+			},
+		},
+	};
+}
diff --git a/packages/integrations/markdoc/src/nodes/heading.ts b/packages/integrations/markdoc/src/heading-ids.ts
index 0210e9b90..57b84d059 100644
--- a/packages/integrations/markdoc/src/nodes/heading.ts
+++ b/packages/integrations/markdoc/src/heading-ids.ts
@@ -1,13 +1,8 @@
 import Markdoc, { type ConfigType, type RenderableTreeNode, type Schema } from '@markdoc/markdoc';
 import Slugger from 'github-slugger';
-import { getTextContent } from '../runtime.js';
-
-type ConfigTypeWithCtx = ConfigType & {
-	// TODO: decide on `ctx` as a convention for config merging
-	ctx: {
-		headingSlugger: Slugger;
-	};
-};
+import { getTextContent } from './runtime.js';
+import type { AstroMarkdocConfig } from './config.js';
+import { MarkdocError } from './utils.js';
 
 function getSlug(
 	attributes: Record<string, any>,
@@ -24,16 +19,31 @@ function getSlug(
 	return slug;
 }
 
+type HeadingIdConfig = AstroMarkdocConfig<{
+	headingSlugger: Slugger;
+}>;
+
+/*
+	Expose standalone node for users to import in their config.
+	Allows users to apply a custom `render: AstroComponent`
+	and spread our default heading attributes.
+*/
 export const heading: Schema = {
 	children: ['inline'],
 	attributes: {
 		id: { type: String },
 		level: { type: Number, required: true, default: 1 },
 	},
-	transform(node, config: ConfigTypeWithCtx) {
+	transform(node, config: HeadingIdConfig) {
 		const { level, ...attributes } = node.transformAttributes(config);
 		const children = node.transformChildren(config);
 
+		if (!config.ctx?.headingSlugger) {
+			throw new MarkdocError({
+				message:
+					'Unexpected problem adding heading IDs to Markdoc file. Did you modify the `ctx.headingSlugger` property in your Markdoc config?',
+			});
+		}
 		const slug = getSlug(attributes, children, config.ctx.headingSlugger);
 
 		const render = config.nodes?.heading?.render ?? `h${level}`;
@@ -49,9 +59,9 @@ export const heading: Schema = {
 	},
 };
 
-export function setupHeadingConfig(): ConfigTypeWithCtx {
+// Called internally to ensure `ctx` is generated per-file, instead of per-build.
+export function setupHeadingConfig(): HeadingIdConfig {
 	const headingSlugger = new Slugger();
-
 	return {
 		ctx: {
 			headingSlugger,
diff --git a/packages/integrations/markdoc/src/index.ts b/packages/integrations/markdoc/src/index.ts
index 627f08c77..64ae4cbc0 100644
--- a/packages/integrations/markdoc/src/index.ts
+++ b/packages/integrations/markdoc/src/index.ts
@@ -52,7 +52,11 @@ export default function markdocIntegration(legacyConfig?: any): AstroIntegration
 					async getRenderModule({ entry, viteId }) {
 						const ast = Markdoc.parse(entry.body);
 						const pluginContext = this;
-						const markdocConfig = setupConfig(userMarkdocConfig, entry);
+						const markdocConfig = setupConfig(
+							userMarkdocConfig,
+							entry,
+							markdocConfigResult?.fileUrl.pathname
+						);
 
 						const validationErrors = Markdoc.validate(ast, markdocConfig).filter((e) => {
 							return (
diff --git a/packages/integrations/markdoc/src/nodes/index.ts b/packages/integrations/markdoc/src/nodes/index.ts
deleted file mode 100644
index 4cd7e3667..000000000
--- a/packages/integrations/markdoc/src/nodes/index.ts
+++ /dev/null
@@ -1,4 +0,0 @@
-import { heading } from './heading.js';
-export { setupHeadingConfig } from './heading.js';
-
-export const nodes = { heading };
diff --git a/packages/integrations/markdoc/src/runtime.ts b/packages/integrations/markdoc/src/runtime.ts
index 3164cda13..4c5614b56 100644
--- a/packages/integrations/markdoc/src/runtime.ts
+++ b/packages/integrations/markdoc/src/runtime.ts
@@ -1,32 +1,56 @@
 import type { MarkdownHeading } from '@astrojs/markdown-remark';
-import Markdoc, {
-	type ConfigType as MarkdocConfig,
-	type RenderableTreeNode,
-} from '@markdoc/markdoc';
+import Markdoc, { type RenderableTreeNode } from '@markdoc/markdoc';
 import type { ContentEntryModule } from 'astro';
-import { setupHeadingConfig } from './nodes/index.js';
+import { setupHeadingConfig } from './heading-ids.js';
+import type { AstroMarkdocConfig } from './config.js';
+import { MarkdocError } from './utils.js';
 
 /** Used to call `Markdoc.transform()` and `Markdoc.Ast` in runtime modules */
 export { default as Markdoc } from '@markdoc/markdoc';
 
 /**
  * Merge user config with default config and set up context (ex. heading ID slugger)
- * Called on each file's individual transform
+ * Called on each file's individual transform.
+ * TODO: virtual module to merge configs per-build instead of per-file?
  */
-export function setupConfig(userConfig: MarkdocConfig, entry: ContentEntryModule): MarkdocConfig {
-	const defaultConfig: MarkdocConfig = {
-		// `setupXConfig()` could become a "plugin" convention as well?
+export function setupConfig(
+	userConfig: AstroMarkdocConfig,
+	entry: ContentEntryModule,
+	markdocConfigPath?: string
+): Omit<AstroMarkdocConfig, 'extends'> {
+	let defaultConfig: AstroMarkdocConfig = {
 		...setupHeadingConfig(),
 		variables: { entry },
 	};
+
+	if (userConfig.extends) {
+		for (const extension of userConfig.extends) {
+			if (extension instanceof Promise) {
+				throw new MarkdocError({
+					message: 'An extension passed to `extends` in your markdoc config returns a Promise.',
+					hint: 'Call `await` for async extensions. Example: `extends: [await myExtension()]`',
+					location: {
+						file: markdocConfigPath,
+					},
+				});
+			}
+
+			defaultConfig = mergeConfig(defaultConfig, extension);
+		}
+	}
+
 	return mergeConfig(defaultConfig, userConfig);
 }
 
 /** Merge function from `@markdoc/markdoc` internals */
-function mergeConfig(configA: MarkdocConfig, configB: MarkdocConfig): MarkdocConfig {
+function mergeConfig(configA: AstroMarkdocConfig, configB: AstroMarkdocConfig): AstroMarkdocConfig {
 	return {
 		...configA,
 		...configB,
+		ctx: {
+			...configA.ctx,
+			...configB.ctx,
+		},
 		tags: {
 			...configA.tags,
 			...configB.tags,
diff --git a/packages/integrations/markdoc/test/syntax-highlighting.test.js b/packages/integrations/markdoc/test/syntax-highlighting.test.js
new file mode 100644
index 000000000..ef1845eb9
--- /dev/null
+++ b/packages/integrations/markdoc/test/syntax-highlighting.test.js
@@ -0,0 +1,89 @@
+import { parseHTML } from 'linkedom';
+import { expect } from 'chai';
+import Markdoc from '@markdoc/markdoc';
+import { shiki } from '../dist/config.js';
+import { setupConfig } from '../dist/runtime.js';
+import { isHTMLString } from 'astro/runtime/server/index.js';
+
+const entry = `
+\`\`\`ts
+const highlighting = true;
+\`\`\`
+
+\`\`\`css
+.highlighting {
+	color: red;
+}
+\`\`\`
+`;
+
+describe('Markdoc - syntax highlighting', () => {
+	it('transforms with defaults', async () => {
+		const ast = Markdoc.parse(entry);
+		const content = Markdoc.transform(ast, await getConfigExtendingShiki());
+
+		expect(content.children).to.have.lengthOf(2);
+		for (const codeBlock of content.children) {
+			expect(isHTMLString(codeBlock)).to.be.true;
+
+			const pre = parsePreTag(codeBlock);
+			expect(pre.classList).to.include('astro-code');
+			expect(pre.classList).to.include('github-dark');
+		}
+	});
+	it('transforms with `theme` property', async () => {
+		const ast = Markdoc.parse(entry);
+		const content = Markdoc.transform(
+			ast,
+			await getConfigExtendingShiki({
+				theme: 'dracula',
+			})
+		);
+		expect(content.children).to.have.lengthOf(2);
+		for (const codeBlock of content.children) {
+			expect(isHTMLString(codeBlock)).to.be.true;
+
+			const pre = parsePreTag(codeBlock);
+			expect(pre.classList).to.include('astro-code');
+			expect(pre.classList).to.include('dracula');
+		}
+	});
+	it('transforms with `wrap` property', async () => {
+		const ast = Markdoc.parse(entry);
+		const content = Markdoc.transform(
+			ast,
+			await getConfigExtendingShiki({
+				wrap: true,
+			})
+		);
+		expect(content.children).to.have.lengthOf(2);
+		for (const codeBlock of content.children) {
+			expect(isHTMLString(codeBlock)).to.be.true;
+
+			const pre = parsePreTag(codeBlock);
+			expect(pre.getAttribute('style')).to.include('white-space: pre-wrap');
+			expect(pre.getAttribute('style')).to.include('word-wrap: break-word');
+		}
+	});
+});
+
+/**
+ * @param {import('astro').ShikiConfig} config
+ * @returns {import('../src/config.js').AstroMarkdocConfig}
+ */
+async function getConfigExtendingShiki(config) {
+	return setupConfig({
+		extends: [await shiki(config)],
+	});
+}
+
+/**
+ * @param {string} html
+ * @returns {HTMLPreElement}
+ */
+function parsePreTag(html) {
+	const { document } = parseHTML(html);
+	const pre = document.querySelector('pre');
+	expect(pre).to.exist;
+	return pre;
+}
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
index f91304cf4..f5f47aa8c 100644
--- a/pnpm-lock.yaml
+++ b/pnpm-lock.yaml
@@ -4025,6 +4025,9 @@ importers:
       rollup:
         specifier: ^3.20.1
         version: 3.20.1
+      shiki:
+        specifier: ^0.14.1
+        version: 0.14.1
       vite:
         specifier: ^4.3.1
         version: 4.3.1(@types/node@18.16.3)(sass@1.52.2)