Add plugin config for custom code highlighting.

Author: koculu

packages/ts-ssg/README.md

@@ -370,6 +370,26 @@ Plugins apply in order:
 
 Each plugin is checked when the build starts. An unknown field or hook name, such as `onPageRender`, a value of the wrong type, two plugins sharing a name, or two plugins defining the same skin, component, template, or generated page fails the build with the plugin's name. For a single hook, pass an inline plugin: `{ name: 'site', hooks: { … } }`.
 
+### Code highlighting
+
+Use `markdown.resolveHighlighter(lang)` to select a code-to-HTML callback for a language:
+
+```ts
+import { definePlugin } from '@purestack/ts-ssg'
+import { renderCustomCode } from './custom-code'
+
+export const customCodePlugin = definePlugin({
+  name: 'custom-code',
+  markdown: {
+    resolveHighlighter: (lang) => lang === 'custom' ? renderCustomCode : null,
+  },
+})
+```
+
+The resolver can be synchronous or asynchronous and receives the trimmed, lowercase language, or `''` when absent. Its callback accepts `(code, lang?)` and returns the HTML as `string | Promise<string>`. Resolvers are awaited in plugin order; the first callback returned handles the block. If all return `null`, the configured default highlighter handles it, including its existing behavior for unknown languages.
+
+Custom callbacks also work when `mdx.disableHighlighter` is `true`. In that case, if all resolvers return `null`, the source is HTML-escaped and displayed in a plain `<pre><code>` block without syntax highlighting.
+
 ### Hooks
 
 Every page render, in a full build or a dev server re-render, runs the page hooks:

packages/ts-ssg/src/build/incremental/index.ts

@@ -97,7 +97,7 @@ async function createIncrementalRuntime(
   const minifyScripts = publishOptions.enabled === true
   const failOnAssetError = publishOptions.enabled === true
   const mdx = {
-    ...(await resolveMdxBuildOptions(config.mdx)),
+    ...(await resolveMdxBuildOptions(config.mdx, plugins)),
     contentProcessor: resolvePluginContentProcessor(plugins),
   }
   themes.setOptions(config.style.theme)

packages/ts-ssg/src/build/incremental/support.ts

@@ -22,7 +22,10 @@ import {
 } from '../../mdx/highlight'
 import { createHljsHighlighter } from '../../mdx/highlightjs'
 import { generatePluginPages } from '../../plugins/generated-pages'
-import type { PureStackPlugin } from '../../plugins/plugin'
+import {
+  type PureStackPlugin,
+  resolvePluginHighlighter,
+} from '../../plugins/plugin'
 import {
   assertUniqueContentRoutes,
   resolveRouteInfo,
@@ -128,8 +131,12 @@ export class ManifestContentIndex {
 
 export async function resolveMdxBuildOptions(
   mdx: Partial<SiteMdxConfig> | undefined,
+  plugins: readonly PureStackPlugin[] = [],
 ): Promise<MdxRenderOptions> {
-  const highlighter = await resolveHighlighter(mdx)
+  const highlighter = resolvePluginHighlighter(
+    plugins,
+    await resolveHighlighter(mdx),
+  )
   return {
     highlighter,
     compileMdAsMdx: mdx?.compileMdAsMdx ?? true,

packages/ts-ssg/src/mdx/codeHighlighting.test.ts

@@ -0,0 +1,113 @@
+import { describe, expect, it } from 'vitest'
+import { createHljsHighlighter } from './highlightjs'
+import { compileMarkdown } from './md'
+import { compileMdx } from './mdx'
+
+describe('code highlighting', () => {
+  it.each([compileMarkdown, compileMdx])(
+    'awaits arbitrary block HTML through the existing hook',
+    async (compile) => {
+      const html =
+        '<figure r-pre><svg viewBox="0 0 10 10"><title>{{ literal }}</title></svg></figure><p>Caption</p>'
+      const result = await compile('```diagram\nsource\n```', {
+        highlighter: { codeToHtml: async () => html },
+      })
+      expect(result.bodyHtml).toBe(html)
+    },
+  )
+
+  it('awaits the same hook inside component markup', async () => {
+    const result = await compileMdx(
+      '<Panel>\n```diagram\nsource\n```\n</Panel>',
+      {
+        highlighter: {
+          codeToHtml: async (code) => `<figure r-pre>${code}</figure>`,
+        },
+      },
+    )
+    expect(result.bodyHtml).toContain(
+      '<Panel>\n<figure r-pre>source</figure>\n</Panel>',
+    )
+    expect(result.bodyHtml).not.toContain('<pre')
+  })
+
+  it('preserves callback order across separate component markup segments', async () => {
+    const calls: string[] = []
+    await compileMdx(
+      [
+        '<Panel>',
+        '```custom',
+        'first',
+        '```',
+        '</Panel>',
+        '',
+        'Between components.',
+        '',
+        '<Panel>',
+        '```custom',
+        'second',
+        '```',
+        '</Panel>',
+      ].join('\n'),
+      {
+        highlighter: {
+          codeToHtml: async (code) => {
+            calls.push(`start:${code}`)
+            await Promise.resolve()
+            calls.push(`end:${code}`)
+            return `<figure>${code}</figure>`
+          },
+        },
+      },
+    )
+    expect(calls).toEqual([
+      'start:first',
+      'end:first',
+      'start:second',
+      'end:second',
+    ])
+  })
+
+  it.each([compileMarkdown, compileMdx])(
+    'honors empty custom HTML without restoring a code block',
+    async (compile) => {
+      const source =
+        compile === compileMdx
+          ? '<Panel>\n```custom\nsource\n```\n</Panel>'
+          : '```custom\nsource\n```'
+      const result = await compile(source, {
+        highlighter: { codeToHtml: () => '' },
+      })
+      expect(result.bodyHtml).not.toContain('source')
+      expect(result.bodyHtml).not.toContain('<pre')
+    },
+  )
+
+  it('preserves built-in block and inline highlighting with an asynchronous wrapper', async () => {
+    const fallback = createHljsHighlighter()
+    const source =
+      '```ts\nconst answer = 42\n```\n\n`value`\n\n<Panel>\n```ts\nconst nested = 1\n```\n`nested`\n</Panel>'
+    const original = await compileMdx(source, { highlighter: fallback })
+    const result = await compileMdx(source, {
+      highlighter: {
+        codeToHtml: async (code, lang) => fallback.codeToHtml(code, lang),
+      },
+    })
+    expect(result).toEqual(original)
+  })
+
+  it.each([compileMarkdown, compileMdx])(
+    'propagates asynchronous highlighter errors',
+    async (compile) => {
+      await expect(
+        compile('```diagram\nbroken\n```', {
+          highlighter: {
+            codeToHtml: async () => {
+              throw new Error('invalid diagram')
+            },
+          },
+        }),
+      ).rejects.toThrow('invalid diagram')
+    },
+  )
+})

…ages/ts-ssg/src/mdx/shikiHighlighting.ts …kages/ts-ssg/src/mdx/codeHighlighting.tspackages/ts-ssg/src/mdx/shikiHighlighting.ts renamed to packages/ts-ssg/src/mdx/codeHighlighting.ts packages/ts-ssg/src/mdx/shikiHighlighting.ts renamed to packages/ts-ssg/src/mdx/codeHighlighting.ts

@@ -1,62 +1,64 @@
 import type { Element, Root, Text } from 'hast'
 import { fromHtml } from 'hast-util-from-html'
 import { toHtml } from 'hast-util-to-html'
+import type { Raw } from 'mdast-util-to-hast'
 
 import type { MdxCodeHighlighter } from './highlight'
 
-export function applyShikiHighlighting(
+/** Applies the configured highlighter to code blocks and inline code. */
+export async function applyCodeHighlighting(
   root: Root,
   highlighter: MdxCodeHighlighter,
 ) {
-  const visit = (node: Root | Element) => {
+  const visit = async (node: Root | Element) => {
     const children = node.children ?? []
     for (let index = 0; index < children.length; index += 1) {
       const child = children[index]
       if (child.type !== 'element') continue
       if (child.tagName === 'pre') {
-        const highlighted = highlightPre(child, highlighter)
+        const highlighted = await highlightCodeBlock(child, highlighter)
         if (highlighted) {
           children[index] = highlighted
           continue
         }
       }
       if (child.tagName === 'code' && !isPre(node)) {
-        const highlighted = highlightInlineCode(child, highlighter)
+        const highlighted = await highlightInlineCode(child, highlighter)
         if (highlighted) {
           children[index] = highlighted
           continue
         }
       }
       if (child.children && child.children.length > 0) {
-        visit(child)
+        await visit(child)
       }
     }
   }
-  visit(root)
+  await visit(root)
 }
 
-function highlightPre(
+async function highlightCodeBlock(
   pre: Element,
   highlighter: MdxCodeHighlighter,
-): Element | undefined {
+): Promise<Element | Raw | undefined> {
   const code = pre.children?.find(
     (child) => child.type === 'element' && child.tagName === 'code',
   ) as Element | undefined
   if (!code) return undefined
   const language = resolveLanguage(code) ?? resolveLanguage(pre)
   const text = collectText(code).trimEnd()
   if (text.length === 0) return undefined
-  return renderHighlightedPre(text, language, highlighter)
+  return renderHighlightedCodeBlock(text, language, highlighter)
 }
 
-function highlightInlineCode(
+async function highlightInlineCode(
   code: Element,
   highlighter: MdxCodeHighlighter,
-): Element | undefined {
+): Promise<Element | undefined> {
   const language = resolveLanguage(code)
   const text = collectText(code).trimEnd()
   if (text.length === 0) return undefined
-  const html = highlighter.codeToHtml(text, language)
+  const html = await highlighter.codeToHtml(text, language)
   const parsed = fromHtml(html, { fragment: true })
   const replacement = parsed.children?.find(
     (child) => child.type === 'element' && child.tagName === 'pre',
@@ -86,38 +88,49 @@ function highlightInlineCode(
   return code
 }
 
-function renderHighlightedPre(
+/** Preserves arbitrary block HTML; a single pre uses the site's code styling. */
+async function renderHighlightedCodeBlock(
   text: string,
   language: string | undefined,
   highlighter: MdxCodeHighlighter,
-): Element | undefined {
-  const html = highlighter.codeToHtml(text, language)
+): Promise<Element | Raw> {
+  const html = await highlighter.codeToHtml(text, language)
   const parsed = fromHtml(html, { fragment: true })
-  const replacement = parsed.children?.find(
-    (child) => child.type === 'element' && child.tagName === 'pre',
-  ) as Element | undefined
-  if (!replacement) return undefined
-  normalizeShikiPre(replacement)
+  const replacement = parsed.children[0]
+  if (
+    parsed.children.length !== 1 ||
+    replacement?.type !== 'element' ||
+    replacement.tagName !== 'pre'
+  ) {
+    return { type: 'raw', value: html }
+  }
+  normalizeHighlightedPre(replacement)
   return replacement
 }
 
-export function renderHighlightedPreHtml(
+/** Renders a code block inside component markup, keeping pre content literal. */
+export async function renderHighlightedCodeBlockHtml(
   text: string,
   language: string | undefined,
   highlighter: MdxCodeHighlighter,
-): string | undefined {
-  const replacement = renderHighlightedPre(text, language, highlighter)
-  if (!replacement) return undefined
+): Promise<string> {
+  const replacement = await renderHighlightedCodeBlock(
+    text,
+    language,
+    highlighter,
+  )
+  if (replacement.type === 'raw') return replacement.value
   replacement.properties['r-pre'] = ''
   return toHtml(replacement, { allowDangerousHtml: true })
 }
 
-export function renderHighlightedInlineCodeHtml(
+/** Extracts inline code from the highlighter's pre/code output. */
+export async function renderHighlightedInlineCodeHtml(
   text: string,
   language: string | undefined,
   highlighter: MdxCodeHighlighter,
-): string | undefined {
-  const html = highlighter.codeToHtml(text, language)
+): Promise<string | undefined> {
+  const html = await highlighter.codeToHtml(text, language)
   const parsed = fromHtml(html, { fragment: true })
   const replacement = parsed.children?.find(
     (child) => child.type === 'element' && child.tagName === 'pre',
@@ -184,7 +197,8 @@ function isPre(node: Root | Element) {
   return node.type === 'element' && node.tagName === 'pre'
 }
 
-function normalizeShikiPre(pre: Element) {
+/** The site's shared code styles use the shiki class for every highlighter. */
+function normalizeHighlightedPre(pre: Element) {
   if (!pre.properties) pre.properties = {}
   pre.properties.className = mergeClassNames(pre.properties.className, 'shiki')
   const style =

packages/ts-ssg/src/mdx/compile.ts

@@ -5,8 +5,8 @@ import type { Root as MdastRoot } from 'mdast'
 import rehypeStringify from 'rehype-stringify'
 import remarkRehype from 'remark-rehype'
 import { type PluggableList, unified } from 'unified'
+import { applyCodeHighlighting } from './codeHighlighting'
 import type { MdxCodeHighlighter } from './highlight'
-import { applyShikiHighlighting } from './shikiHighlighting'
 
 export interface MdxCompileResult {
   bodyHtml: string
@@ -56,7 +56,7 @@ export async function compileAstToHtml(
   wrapTablesInScrollContainers(tree)
   const outline = collectOutline(tree)
   if (options.highlighter) {
-    applyShikiHighlighting(tree, options.highlighter)
+    await applyCodeHighlighting(tree, options.highlighter)
   }
   keepCodeLiteral(tree)
   const bodyHtml = String(

packages/ts-ssg/src/mdx/highlight.ts

@@ -1,7 +1,7 @@
 import { createHighlighter, type Highlighter } from 'shiki'
 
 export interface MdxCodeHighlighter {
-  codeToHtml: (code: string, lang?: string) => string
+  codeToHtml: (code: string, lang?: string) => string | Promise<string>
 }
 
 export interface MdxCodeThemes {

packages/ts-ssg/src/mdx/mdx.test.ts

@@ -19,7 +19,7 @@ describe('createMdxHighlighter', () => {
         // biome-ignore lint/suspicious/noTemplateCurlyInString: Literal TypeScript source for the highlighter.
         'const view = html`<div class="greeting">Hello ${name}</div>`',
       ].join('\n')
-      const html = highlighter.codeToHtml(source, lang)
+      const html = await highlighter.codeToHtml(source, lang)
       const cleanup = ensureDomGlobals()
       try {
         const root = document.createElement('div')

packages/ts-ssg/src/mdx/mdx.ts

@@ -24,7 +24,7 @@ export {
 
 export type { MdxCompileResult, MdxRenderOptions }
 
-export function compileMdx(
+export async function compileMdx(
   source: string,
   options: MdxRenderOptions = {},
 ): Promise<MdxCompileResult> {
@@ -35,7 +35,7 @@ export function compileMdx(
   const file = parser.parse(masked.source)
   restoreRegorMarkup(
     file,
-    renderRegorMarkupCodeFences(masked.segments, options.highlighter),
+    await renderRegorMarkupCodeFences(masked.segments, options.highlighter),
   )
   normalizeMarkupParagraphs(file)
   return compileAstToHtml(file, options)