2 files changed, 100 insertions, 18 deletions
diff --git a/packages/integrations/mdx/src/index.ts b/packages/integrations/mdx/src/index.ts
index 66ab7b837..2ac6cc66a 100644
--- a/packages/integrations/mdx/src/index.ts
+++ b/packages/integrations/mdx/src/index.ts
@@ -1,11 +1,15 @@
-import mdxPlugin, { Options as MdxRollupPluginOptions } from '@mdx-js/rollup';
+import type { RemarkMdxFrontmatterOptions } from 'remark-mdx-frontmatter';
 import type { AstroIntegration } from 'astro';
+import remarkShikiTwoslash from 'remark-shiki-twoslash';
+import { nodeTypes } from '@mdx-js/mdx';
+import rehypeRaw from 'rehype-raw';
+import mdxPlugin, { Options as MdxRollupPluginOptions } from '@mdx-js/rollup';
 import { parse as parseESM } from 'es-module-lexer';
 import remarkFrontmatter from 'remark-frontmatter';
 import remarkGfm from 'remark-gfm';
-import type { RemarkMdxFrontmatterOptions } from 'remark-mdx-frontmatter';
 import remarkMdxFrontmatter from 'remark-mdx-frontmatter';
 import remarkSmartypants from 'remark-smartypants';
+import remarkPrism from './remark-prism.js';
 import { getFileInfo } from './utils.js';
 
 type WithExtends<T> = T | { extends: T };
@@ -23,7 +27,10 @@ type MdxOptions = {
 
 const DEFAULT_REMARK_PLUGINS = [remarkGfm, remarkSmartypants];
 
-function handleExtends<T>(config: WithExtends<T[] | undefined>, defaults: T[] = []): T[] {
+function handleExtends<T>(
+	config: WithExtends<T[] | undefined>,
+	defaults: T[] = [],
+): T[] {
 	if (Array.isArray(config)) return config;
 
 	return [...defaults, ...(config?.extends ?? [])];
@@ -35,27 +42,43 @@ export default function mdx(mdxOptions: MdxOptions = {}): AstroIntegration {
 		hooks: {
 			'astro:config:setup': ({ updateConfig, config, addPageExtension, command }: any) => {
 				addPageExtension('.mdx');
+				let remarkPlugins = handleExtends(mdxOptions.remarkPlugins, DEFAULT_REMARK_PLUGINS);
+				let rehypePlugins = handleExtends(mdxOptions.rehypePlugins);
+
+				if (config.markdown.syntaxHighlight === 'shiki') {
+					remarkPlugins.push([
+						// Default export still requires ".default" chaining for some reason
+						// Workarounds tried:
+						// - "import * as remarkShikiTwoslash"
+						// - "import { default as remarkShikiTwoslash }"
+						(remarkShikiTwoslash as any).default,
+						config.markdown.shikiConfig,
+					]);
+					rehypePlugins.push([rehypeRaw, { passThrough: nodeTypes }]);
+				}
+				
+				if (config.markdown.syntaxHighlight === 'prism') {
+					remarkPlugins.push(remarkPrism);
+					rehypePlugins.push([rehypeRaw, { passThrough: nodeTypes }]);
+				}
+
+				remarkPlugins.push(remarkFrontmatter);
+				remarkPlugins.push([
+					remarkMdxFrontmatter,
+					{
+						name: 'frontmatter',
+						...mdxOptions.frontmatterOptions,
+					},
+				]);
+
 				updateConfig({
 					vite: {
 						plugins: [
 							{
 								enforce: 'pre',
 								...mdxPlugin({
-									remarkPlugins: [
-										...handleExtends(mdxOptions.remarkPlugins, DEFAULT_REMARK_PLUGINS),
-										// Frontmatter plugins should always be applied!
-										// We can revisit this if a strong use case to *remove*
-										// YAML frontmatter via config is reported.
-										remarkFrontmatter,
-										[
-											remarkMdxFrontmatter,
-											{
-												name: 'frontmatter',
-												...mdxOptions.frontmatterOptions,
-											},
-										],
-									],
-									rehypePlugins: handleExtends(mdxOptions.rehypePlugins),
+									remarkPlugins,
+									rehypePlugins,
 									jsx: true,
 									jsxImportSource: 'astro',
 									// Note: disable `.md` support
diff --git a/packages/integrations/mdx/src/remark-prism.ts b/packages/integrations/mdx/src/remark-prism.ts
new file mode 100644
index 000000000..019c3984b
--- /dev/null
+++ b/packages/integrations/mdx/src/remark-prism.ts
@@ -0,0 +1,59 @@
+// TODO: discuss extracting this file to @astrojs/prism
+import { addAstro } from '@astrojs/prism/internal';
+import Prism from 'prismjs';
+import loadLanguages from 'prismjs/components/index.js';
+import { visit } from 'unist-util-visit';
+
+const languageMap = new Map([['ts', 'typescript']]);
+
+function runHighlighter(lang: string, code: string) {
+	let classLanguage = `language-${lang}`;
+
+	if (lang == null) {
+		lang = 'plaintext';
+	}
+
+	const ensureLoaded = (language: string) => {
+		if (language && !Prism.languages[language]) {
+			loadLanguages([language]);
+		}
+	};
+
+	if (languageMap.has(lang)) {
+		ensureLoaded(languageMap.get(lang)!);
+	} else if (lang === 'astro') {
+		ensureLoaded('typescript');
+		addAstro(Prism);
+	} else {
+		ensureLoaded('markup-templating'); // Prism expects this to exist for a number of other langs
+		ensureLoaded(lang);
+	}
+
+	if (lang && !Prism.languages[lang]) {
+		// eslint-disable-next-line no-console
+		console.warn(`Unable to load the language: ${lang}`);
+	}
+
+	const grammar = Prism.languages[lang];
+	let html = code;
+	if (grammar) {
+		html = Prism.highlight(code, grammar, lang);
+	}
+
+	return { classLanguage, html };
+}
+
+/**  */
+export default function remarkPrism() {
+	return (tree: any) => visit(tree, 'code', (node: any) => {
+		let { lang, value } = node;
+		node.type = 'html';
+
+		let { html, classLanguage } = runHighlighter(lang, value);
+		let classes = [classLanguage];
+		node.value = `<pre class="${classes.join(
+			' '
+		)}"><code class="${classLanguage}">${html}</code></pre>`;
+		return node;
+	});
+}