[MDX] Switch from Shiki Twoslash -> Astro Markdown highlighter (#4292)

* freat: twoslash -> Astro shiki parser

* test: update shiki style check

* feat: always apply rehypeRaw

* deps: move remark-shiki-twoslash to dev

* test: add shiki-twoslash test

* docs: update readme with twoslash example

* chore: changeset

* nit: remove "describe('disabled')"

5 files changed, 138 insertions, 19 deletions

diff --git a/packages/integrations/mdx/README.md b/packages/integrations/mdx/README.md
index a77733589..c589673d7 100644
--- a/packages/integrations/mdx/README.md
+++ b/packages/integrations/mdx/README.md
@@ -253,7 +253,7 @@ const { title, fancyJsHelper } = Astro.props;
 
 The MDX integration respects [your project's `markdown.syntaxHighlight` configuration](https://docs.astro.build/en/guides/markdown-content/#syntax-highlighting).
 
-We will highlight your code blocks with [Shiki](https://github.com/shikijs/shiki) by default [using Shiki twoslash](https://shikijs.github.io/twoslash/). You can customize [this remark plugin](https://www.npmjs.com/package/remark-shiki-twoslash) using the `markdown.shikiConfig` option in your `astro.config`. For example, you can apply a different built-in theme like so:
+We will highlight your code blocks with [Shiki](https://github.com/shikijs/shiki) by default. You can customize this highlighter using the `markdown.shikiConfig` option in your `astro.config`. For example, you can apply a different built-in theme like so:
 
 ```js
 // astro.config.mjs
@@ -285,6 +285,23 @@ export default {
 
 This applies a minimal Prism renderer with added support for `astro` code blocks. Visit [our "Prism configuration" docs](https://docs.astro.build/en/guides/markdown-content/#prism-configuration) for more on using Prism with Astro.
 
+#### Switch to a custom syntax highlighter
+
+You may want to apply your own syntax highlighter too. If your highlighter offers a remark or rehype plugin, you can flip off our syntax highlighting by setting `markdown.syntaxHighlight: false` and wiring up your plugin. For example, say you want to apply [Shiki Twoslash's remark plugin](https://www.npmjs.com/package/remark-shiki-twoslash):
+
+```js
+// astro.config.mjs
+import shikiTwoslash from 'remark-shiki-twoslash';
+
+export default {
+  markdown: {
+  syntaxHighlight: false,
+  },
+  integrations: [mdx({
+    remarkPlugins: [shikiTwoslash, { /* Shiki Twoslash config */ }],
+  })],
+```
+
 ## Configuration
 
 ### remarkPlugins
diff --git a/packages/integrations/mdx/package.json b/packages/integrations/mdx/package.json
index 618eaabfe..636cdb643 100644
--- a/packages/integrations/mdx/package.json
+++ b/packages/integrations/mdx/package.json
@@ -39,7 +39,6 @@
     "rehype-raw": "^6.1.1",
     "remark-frontmatter": "^4.0.1",
     "remark-gfm": "^3.0.1",
-    "remark-shiki-twoslash": "^3.1.0",
     "remark-smartypants": "^2.0.0",
     "shiki": "^0.10.1",
     "unist-util-visit": "^4.1.0",
@@ -56,6 +55,7 @@
     "mdast-util-to-string": "^3.1.0",
     "mocha": "^9.2.2",
     "reading-time": "^1.5.0",
+    "remark-shiki-twoslash": "^3.1.0",
     "remark-toc": "^8.0.1"
   },
   "engines": {
diff --git a/packages/integrations/mdx/src/index.ts b/packages/integrations/mdx/src/index.ts
index 17fe0cd74..72fbbeb6c 100644
--- a/packages/integrations/mdx/src/index.ts
+++ b/packages/integrations/mdx/src/index.ts
@@ -4,13 +4,13 @@ import type { AstroConfig, AstroIntegration } from 'astro';
 import { parse as parseESM } from 'es-module-lexer';
 import rehypeRaw from 'rehype-raw';
 import remarkGfm from 'remark-gfm';
-import remarkShikiTwoslash from 'remark-shiki-twoslash';
 import remarkSmartypants from 'remark-smartypants';
 import { VFile } from 'vfile';
 import type { Plugin as VitePlugin } from 'vite';
 import { rehypeApplyFrontmatterExport, remarkInitializeAstroData } from './astro-data-utils.js';
 import rehypeCollectHeadings from './rehype-collect-headings.js';
 import remarkPrism from './remark-prism.js';
+import remarkShiki from './remark-shiki.js';
 import { getFileInfo, parseFrontmatter } from './utils.js';
 
 type WithExtends<T> = T | { extends: T };
@@ -38,22 +38,17 @@ function handleExtends<T>(config: WithExtends<T[] | undefined>, defaults: T[] =
 	return [...defaults, ...(config?.extends ?? [])];
 }
 
-function getRemarkPlugins(
+async function getRemarkPlugins(
 	mdxOptions: MdxOptions,
 	config: AstroConfig
-): MdxRollupPluginOptions['remarkPlugins'] {
+): Promise<MdxRollupPluginOptions['remarkPlugins']> {
 	let remarkPlugins = [
 		// Initialize vfile.data.astroExports before all plugins are run
 		remarkInitializeAstroData,
 		...handleExtends(mdxOptions.remarkPlugins, DEFAULT_REMARK_PLUGINS),
 	];
 	if (config.markdown.syntaxHighlight === 'shiki') {
-		// Default export still requires ".default" chaining for some reason
-		// Workarounds tried:
-		// - "import * as remarkShikiTwoslash"
-		// - "import { default as remarkShikiTwoslash }"
-		const shikiTwoslash = (remarkShikiTwoslash as any).default ?? remarkShikiTwoslash;
-		remarkPlugins.push([shikiTwoslash, config.markdown.shikiConfig]);
+		remarkPlugins.push([await remarkShiki(config.markdown.shikiConfig)]);
 	}
 	if (config.markdown.syntaxHighlight === 'prism') {
 		remarkPlugins.push(remarkPrism);
@@ -65,11 +60,11 @@ function getRehypePlugins(
 	mdxOptions: MdxOptions,
 	config: AstroConfig
 ): MdxRollupPluginOptions['rehypePlugins'] {
-	let rehypePlugins = handleExtends(mdxOptions.rehypePlugins, DEFAULT_REHYPE_PLUGINS);
+	let rehypePlugins = [
+		[rehypeRaw, { passThrough: nodeTypes }] as any,
+		...handleExtends(mdxOptions.rehypePlugins, DEFAULT_REHYPE_PLUGINS),
+	];
 
-	if (config.markdown.syntaxHighlight === 'shiki' || config.markdown.syntaxHighlight === 'prism') {
-		rehypePlugins.unshift([rehypeRaw, { passThrough: nodeTypes }]);
-	}
 	// getHeadings() is guaranteed by TS, so we can't allow user to override
 	rehypePlugins.unshift(rehypeCollectHeadings);
 
@@ -80,11 +75,11 @@ export default function mdx(mdxOptions: MdxOptions = {}): AstroIntegration {
 	return {
 		name: '@astrojs/mdx',
 		hooks: {
-			'astro:config:setup': ({ updateConfig, config, addPageExtension, command }: any) => {
+			'astro:config:setup': async ({ updateConfig, config, addPageExtension, command }: any) => {
 				addPageExtension('.mdx');
 
 				const mdxPluginOpts: MdxRollupPluginOptions = {
-					remarkPlugins: getRemarkPlugins(mdxOptions, config),
+					remarkPlugins: await getRemarkPlugins(mdxOptions, config),
 					rehypePlugins: getRehypePlugins(mdxOptions, config),
 					jsx: true,
 					jsxImportSource: 'astro',
diff --git a/packages/integrations/mdx/src/remark-shiki.ts b/packages/integrations/mdx/src/remark-shiki.ts
new file mode 100644
index 000000000..76a1d275f
--- /dev/null
+++ b/packages/integrations/mdx/src/remark-shiki.ts
@@ -0,0 +1,85 @@
+import type { ShikiConfig } from 'astro';
+import type * as shiki from 'shiki';
+import { getHighlighter } from 'shiki';
+import { visit } from 'unist-util-visit';
+
+/**
+ * getHighlighter() is the most expensive step of Shiki. Instead of calling it on every page,
+ * cache it here as much as possible. Make sure that your highlighters can be cached, state-free.
+ * We make this async, so that multiple calls to parse markdown still share the same highlighter.
+ */
+const highlighterCacheAsync = new Map<string, Promise<shiki.Highlighter>>();
+
+const remarkShiki = async ({ langs = [], theme = 'github-dark', wrap = false }: ShikiConfig) => {
+	const cacheID: string = typeof theme === 'string' ? theme : theme.name;
+	let highlighterAsync = highlighterCacheAsync.get(cacheID);
+	if (!highlighterAsync) {
+		highlighterAsync = getHighlighter({ theme });
+		highlighterCacheAsync.set(cacheID, highlighterAsync);
+	}
+	const highlighter = await highlighterAsync;
+
+	// NOTE: There may be a performance issue here for large sites that use `lang`.
+	// Since this will be called on every page load. Unclear how to fix this.
+	for (const lang of langs) {
+		await highlighter.loadLanguage(lang);
+	}
+
+	return () => (tree: any) => {
+		visit(tree, 'code', (node) => {
+			let lang: string;
+
+			if (typeof node.lang === 'string') {
+				const langExists = highlighter.getLoadedLanguages().includes(node.lang);
+				if (langExists) {
+					lang = node.lang;
+				} else {
+					// eslint-disable-next-line no-console
+					console.warn(`The language "${node.lang}" doesn't exist, falling back to plaintext.`);
+					lang = 'plaintext';
+				}
+			} else {
+				lang = 'plaintext';
+			}
+
+			let html = highlighter!.codeToHtml(node.value, { lang });
+
+			// Q: Couldn't 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;
+
+			// Replace "shiki" class naming with "astro".
+			html = html.replace('<pre class="shiki"', `<pre class="astro-code"`);
+			// Replace "shiki" css variable naming with "astro".
+			html = html.replace(
+				/style="(background-)?color: var\(--shiki-/g,
+				'style="$1color: var(--astro-code-'
+			);
+			// Add "user-select: none;" for "+"/"-" diff symbols
+			if (node.lang === 'diff') {
+				html = html.replace(
+					/<span class="line"><span style="(.*?)">([\+|\-])/g,
+					'<span class="line"><span style="$1"><span style="user-select: none;">$2</span>'
+				);
+			}
+			// Handle code wrapping
+			// if wrap=null, do nothing.
+			if (wrap === false) {
+				html = html.replace(/style="(.*?)"/, 'style="$1; overflow-x: auto;"');
+			} else if (wrap === true) {
+				html = html.replace(
+					/style="(.*?)"/,
+					'style="$1; overflow-x: auto; white-space: pre-wrap; word-wrap: break-word;"'
+				);
+			}
+
+			node.type = 'html';
+			node.value = html;
+			node.children = [];
+		});
+	};
+};
+
+export default remarkShiki;
diff --git a/packages/integrations/mdx/test/mdx-syntax-highlighting.test.js b/packages/integrations/mdx/test/mdx-syntax-highlighting.test.js
index 5544aef56..fc81a7fe3 100644
--- a/packages/integrations/mdx/test/mdx-syntax-highlighting.test.js
+++ b/packages/integrations/mdx/test/mdx-syntax-highlighting.test.js
@@ -3,6 +3,7 @@ import mdx from '@astrojs/mdx';
 import { expect } from 'chai';
 import { parseHTML } from 'linkedom';
 import { loadFixture } from '../../../astro/test/test-utils.js';
+import shikiTwoslash from 'remark-shiki-twoslash';
 
 const FIXTURE_ROOT = new URL('./fixtures/mdx-syntax-hightlighting/', import.meta.url);
 
@@ -21,8 +22,9 @@ describe('MDX syntax highlighting', () => {
 			const html = await fixture.readFile('/index.html');
 			const { document } = parseHTML(html);
 
-			const shikiCodeBlock = document.querySelector('pre.shiki');
+			const shikiCodeBlock = document.querySelector('pre.astro-code');
 			expect(shikiCodeBlock).to.not.be.null;
+			expect(shikiCodeBlock.getAttribute('style')).to.contain('background-color:#0d1117');
 		});
 
 		it('respects markdown.shikiConfig.theme', async () => {
@@ -41,8 +43,9 @@ describe('MDX syntax highlighting', () => {
 			const html = await fixture.readFile('/index.html');
 			const { document } = parseHTML(html);
 
-			const shikiCodeBlock = document.querySelector('pre.shiki.dracula');
+			const shikiCodeBlock = document.querySelector('pre.astro-code');
 			expect(shikiCodeBlock).to.not.be.null;
+			expect(shikiCodeBlock.getAttribute('style')).to.contain('background-color:#282A36');
 		});
 	});
 
@@ -64,4 +67,23 @@ describe('MDX syntax highlighting', () => {
 			expect(prismCodeBlock).to.not.be.null;
 		});
 	});
+
+	it('supports custom highlighter - shiki-twoslash', async () => {
+		const fixture = await loadFixture({
+			root: FIXTURE_ROOT,
+			markdown: {
+				syntaxHighlight: false,
+			},
+			integrations: [mdx({
+				remarkPlugins: [shikiTwoslash.default ?? shikiTwoslash],
+			})],
+		});
+		await fixture.build();
+
+		const html = await fixture.readFile('/index.html');
+		const { document } = parseHTML(html);
+
+		const twoslashCodeBlock = document.querySelector('pre.shiki');
+		expect(twoslashCodeBlock).to.not.be.null;
+	});
 });