Replace Cheerio in mini-TOC with AST-based rehype plugin (#60548)

heiskr · Copilot · web-flow · commit 09b3057d4004 · 2026-04-01T09:35:38.000-07:00
Co-authored-by: Copilot &lt;223556219+Copilot@users.noreply.github.com&gt;
diff --git a/src/content-render/tests/collect-mini-toc.ts b/src/content-render/tests/collect-mini-toc.ts
@@ -0,0 +1,69 @@
+import { describe, expect, test } from 'vitest'
+
+import { renderContent } from '@/content-render/index'
+import type { CollectedHeading } from '@/frame/lib/get-mini-toc-items'
+import type { Context } from '@/types'
+
+describe('collect-mini-toc rehype plugin', () => {
+  test('collects h2 and h3 headings with href and title', async () => {
+    const collectMiniToc: CollectedHeading[] = []
+    await renderContent('## Section one\n\n### Subsection\n\n## Section two', {
+      collectMiniToc,
+    } as Context)
+    expect(collectMiniToc).toHaveLength(3)
+    expect(collectMiniToc[0]).toMatchObject({
+      title: 'Section one',
+      href: '#section-one',
+      headingLevel: 2,
+    })
+    expect(collectMiniToc[1]).toMatchObject({
+      title: 'Subsection',
+      headingLevel: 3,
+    })
+    expect(collectMiniToc[2]).toMatchObject({
+      title: 'Section two',
+      headingLevel: 2,
+    })
+  })
+
+  test('skips headings inside hidden ancestors', async () => {
+    const collectMiniToc: CollectedHeading[] = []
+    const md = [
+      '## Visible heading',
+      '',
+      '<div hidden>',
+      '',
+      '## Hidden heading',
+      '',
+      '</div>',
+      '',
+      '## Another visible',
+    ].join('\n')
+    await renderContent(md, { collectMiniToc } as Context)
+    const titles = collectMiniToc.map((h) => h.title)
+    expect(titles).toContain('Visible heading')
+    expect(titles).toContain('Another visible')
+    expect(titles).not.toContain('Hidden heading')
+  })
+
+  test('detects ghd-tool platform class', async () => {
+    const collectMiniToc: CollectedHeading[] = []
+    const md = ['<div class="ghd-tool mac">', '', '## Mac instructions', '', '</div>'].join('\n')
+    await renderContent(md, { collectMiniToc } as Context)
+    expect(collectMiniToc).toHaveLength(1)
+    expect(collectMiniToc[0].platform).toBe('ghd-tool mac')
+  })
+
+  test('returns empty platform when no ghd-tool wrapper', async () => {
+    const collectMiniToc: CollectedHeading[] = []
+    await renderContent('## Plain heading', { collectMiniToc } as Context)
+    expect(collectMiniToc).toHaveLength(1)
+    expect(collectMiniToc[0].platform).toBe('')
+  })
+
+  test('does not collect when collectMiniToc is not provided', async () => {
+    // Should not throw — plugin is a no-op without collectInto
+    const result = await renderContent('## Heading')
+    expect(result).toContain('Heading')
+  })
+})
diff --git a/src/content-render/unified/collect-mini-toc.ts b/src/content-render/unified/collect-mini-toc.ts
@@ -0,0 +1,81 @@
+import { visitParents } from 'unist-util-visit-parents'
+import { toString } from 'hast-util-to-string'
+import type { Plugin } from 'unified'
+import type { Root, Element, ElementContent } from 'hast'
+import type { CollectedHeading } from '@/frame/lib/get-mini-toc-items'
+
+interface CollectMiniTocOptions {
+  collectInto?: CollectedHeading[]
+}
+
+function hasClassName(el: Element, name: string): boolean {
+  const cls = el.properties?.className
+  if (Array.isArray(cls)) return cls.includes(name)
+  if (typeof cls === 'string') return cls.split(/\s+/).includes(name)
+  return false
+}
+
+function getClassString(el: Element): string {
+  const cls = el.properties?.className
+  if (Array.isArray(cls)) return cls.join(' ')
+  if (typeof cls === 'string') return cls
+  return ''
+}
+
+// Rehype plugin that collects heading data (href, title, level, platform)
+// during rendering, so callers don't need to re-parse the HTML.
+// Place this after heading-links in the processor chain.
+const collectMiniToc: Plugin<[CollectMiniTocOptions], Root> = ({ collectInto }) => {
+  if (!collectInto) return
+
+  return (tree: Root) => {
+    visitParents(tree, 'element', (node, ancestors) => {
+      const el = node as Element
+      if (!/^h[1-6]$/.test(el.tagName)) return
+      if (!el.properties?.id) return
+
+      // Skip headings inside hidden ancestors
+      for (const anc of ancestors) {
+        if (anc.type === 'element') {
+          const ancEl = anc as Element
+          if (ancEl.properties?.hidden === true) return
+        }
+      }
+
+      const headingLevel = parseInt(el.tagName.charAt(1), 10)
+
+      // Find the anchor child that heading-links.ts created
+      const anchor = el.children.find(
+        (child): child is Element =>
+          child.type === 'element' && child.tagName === 'a' && hasClassName(child, 'heading-link'),
+      )
+      if (!anchor) return
+
+      const href = anchor.properties?.href as string | undefined
+      if (!href) return
+
+      // Filter out direct-child <span> elements (and their content).
+      // heading-links.ts always inserts the heading-link-symbol span as a
+      // direct child of the anchor, so filtering direct children is sufficient.
+      const textChildren = (anchor.children || []).filter(
+        (child: ElementContent) =>
+          !(child.type === 'element' && (child as Element).tagName === 'span'),
+      )
+
+      const title = textChildren.map((child) => toString(child)).join('')
+
+      // Detect platform from ghd-tool ancestor
+      let platform = ''
+      for (const anc of ancestors) {
+        if (anc.type === 'element' && hasClassName(anc as Element, 'ghd-tool')) {
+          platform = getClassString(anc as Element)
+          break
+        }
+      }
+
+      collectInto.push({ href, title: title.trim(), headingLevel, platform })
+    })
+  }
+}
+
+export default collectMiniToc
diff --git a/src/content-render/unified/processor.ts b/src/content-render/unified/processor.ts
@@ -20,6 +20,7 @@ import rewriteImgSources from './rewrite-asset-urls'
 import rewriteAssetImgTags from './rewrite-asset-img-tags'
 import useEnglishHeadings from './use-english-headings'
 import headingLinks from './heading-links'
+import collectMiniToc from './collect-mini-toc'
 import rewriteTheadThScope from './rewrite-thead-th-scope'
 import rewriteEmptyTableRows from './rewrite-empty-table-rows'
 import rewriteForRowheaders from './rewrite-for-rowheaders'
@@ -31,6 +32,7 @@ import alerts from './alerts'
 import removeHtmlComments from 'remark-remove-comments'
 import remarkStringify from 'remark-stringify'
 import type { Context, UnifiedProcessor } from '@/content-render/types'
+import type { CollectedHeading } from '@/frame/lib/get-mini-toc-items'
 
 export function createProcessor(context: Context): UnifiedProcessor {
   return (
@@ -74,6 +76,9 @@ export function createProcessor(context: Context): UnifiedProcessor {
         },
       })
       .use(raw)
+      .use(collectMiniToc, {
+        collectInto: context.collectMiniToc as CollectedHeading[] | undefined,
+      })
       .use(wrapProceduralImages)
       .use(rewriteEmptyTableRows)
       .use(rewriteTheadThScope)
diff --git a/src/frame/lib/get-mini-toc-items.ts b/src/frame/lib/get-mini-toc-items.ts
@@ -1,10 +1,13 @@
-import { load } from 'cheerio'
-import type { Element } from 'domhandler'
-import { range } from 'lodash-es'
-
 import { renderContent } from '@/content-render/index'
 import type { Context } from '@/types'
 
+export interface CollectedHeading {
+  href: string
+  title: string
+  headingLevel: number
+  platform: string
+}
+
 interface MiniTocContents {
   href: string
   title: string
@@ -24,81 +27,38 @@ interface FlatTocItem {
   items?: FlatTocItem[]
 }
 
-// Keep maxHeadingLevel=2 for accessibility reasons, see docs-engineering#2701 for more info
-export default function getMiniTocItems(
-  html: string,
+// Build MiniTocItems from pre-collected heading data (from the collect-mini-toc
+// rehype plugin). This is the only path for generating mini-TOC items — headings
+// are collected directly from the AST during rendering, avoiding any HTML
+// re-parsing.
+// Keep maxHeadingLevel=2 for accessibility reasons, see docs-engineering#2701
+export function buildMiniTocFromCollected(
+  collected: CollectedHeading[],
   maxHeadingLevel = 2,
-  headingScope = '',
 ): MiniTocItem[] {
-  const $ = load(html, { xmlMode: true })
-
-  // eg `h2, h3` or `h2, h3, h4` depending on maxHeadingLevel
-  const selector = range(2, maxHeadingLevel + 1)
-    .map((num) => `${headingScope} h${num}`)
-    .join(', ')
-  const headings = $(selector)
-
-  // return an array of objects containing each heading's contents, level, and optional platform.
-  // Article layout uses these as follows:
-  //  - `title` and `link` to render the mini TOC headings
-  //  - `headingLevel` the `2` in `h2`; used for determining required indentation
-  //  - `platform` to show or hide platform-specific headings via client JS
+  const effectiveMax = maxHeadingLevel > 0 ? maxHeadingLevel : 2
+  const headings = collected.filter((h) => h.headingLevel >= 2 && h.headingLevel <= effectiveMax)
 
-  // H1 = highest importance, H6 = lowest importance
   let mostImportantHeadingLevel: number | undefined
-  const flatToc = headings
-    .get()
-    .filter((item) => {
-      const parent = item.parent as Element | null
-      if (!parent || !parent.attribs) return true
-      const { attribs } = parent
-      return !('hidden' in attribs)
-    })
-    .map((item) => {
-      // remove any <span> tags including their content
-      $('span', item).remove()
-
-      // Capture the anchor tag nested within the header, get its href and remove it
-      const anchor = $('a.heading-link', item)
-      const href = anchor.attr('href')
-      if (!href) {
-        // Can happen if the, for example, `<h2>` tag was put there
-        // manually with HTML into the Markdown content. Then it wouldn't
-        // be rendered with an expected `<a class="heading-link" href="#..."`
-        // link in front of it.
-        // The `return null` will be filtered after the `.map()`
-        return null
-      }
-
-      // remove any <strong> tags but leave content
-      $('strong', item).map((i, el) => $(el).replaceWith($(el).contents()))
 
-      const contents: MiniTocContents = { href, title: $(item).text().trim() }
-      const element = $(item)[0] as Element
-      const headingLevel = parseInt(element.name.match(/\d+/)![0], 10) || 0 // the `2` from `h2`
-
-      const platform = $(item).parent('.ghd-tool').attr('class') || ''
-
-      // track the most important heading level while we're looping through the items
-      if (headingLevel < mostImportantHeadingLevel! || mostImportantHeadingLevel === undefined) {
-        mostImportantHeadingLevel = headingLevel
-      }
+  const flatToc: FlatTocItem[] = headings.map((h) => {
+    if (mostImportantHeadingLevel === undefined || h.headingLevel < mostImportantHeadingLevel) {
+      mostImportantHeadingLevel = h.headingLevel
+    }
+    return {
+      contents: { href: h.href, title: h.title },
+      headingLevel: h.headingLevel,
+      platform: h.platform,
+      indentationLevel: 0,
+    }
+  })
 
-      return { contents, headingLevel, platform }
-    })
-    .filter(Boolean)
-    .map((item) => {
-      // set the indentation level for each item based on the most important
-      // heading level in the current article
-      return {
-        ...item!,
-        indentationLevel: item!.headingLevel - mostImportantHeadingLevel!,
-      }
-    })
+  // Set indentation relative to the most important heading
+  for (const item of flatToc) {
+    item.indentationLevel = item.headingLevel - (mostImportantHeadingLevel ?? item.headingLevel)
+  }
 
-  // convert the flatToc to a nested structure to simplify semantic rendering on the client
   const nestedToc = buildNestedToc(flatToc)
-
   return minimalMiniToc(nestedToc)
 }
 
@@ -179,6 +139,10 @@ export async function getAutomatedPageMiniTocItems(
       })
       .join('')
 
-  const toc = await renderContent(titles, context)
-  return getMiniTocItems(toc, depth, '')
+  // Collect headings during render via the rehype plugin
+  const collectMiniToc: CollectedHeading[] = []
+  const renderContext = { ...context, collectMiniToc }
+  await renderContent(titles, renderContext)
+
+  return buildMiniTocFromCollected(collectMiniToc, depth)
 }
diff --git a/src/frame/middleware/render-page.ts b/src/frame/middleware/render-page.ts
@@ -3,7 +3,7 @@ import type { Response } from 'express'
 import type { Failbot } from '@github/failbot'
 import { get } from 'lodash-es'
 
-import getMiniTocItems from '@/frame/lib/get-mini-toc-items'
+import { buildMiniTocFromCollected, type CollectedHeading } from '@/frame/lib/get-mini-toc-items'
 import patterns from '@/frame/lib/patterns'
 import FailBot from '@/observability/lib/failbot'
 import statsd from '@/observability/lib/statsd'
@@ -24,6 +24,13 @@ async function buildRenderedPage(req: ExtendedRequest): Promise<string> {
   if (!page) throw new Error('page not set in context')
   const path = req.pagePath || req.path
 
+  // Set up collection array for the collect-mini-toc rehype plugin only when
+  // the page actually needs a mini-TOC, avoiding unnecessary work.
+  if (page.showMiniToc) {
+    const collectMiniToc: CollectedHeading[] = []
+    context.collectMiniToc = collectMiniToc
+  }
+
   const pageRenderTimed = statsd.asyncTimer(page.render, STATSD_KEY_RENDER, [`path:${path}`])
 
   return (await pageRenderTimed(context)) as string
@@ -39,7 +46,11 @@ function buildMiniTocItems(req: ExtendedRequest) {
     return
   }
 
-  return getMiniTocItems(context.renderedPage || '', 0)
+  // Use headings collected during rendering via the collect-mini-toc rehype plugin.
+  const collected = context.collectMiniToc as CollectedHeading[] | undefined
+  if (collected) {
+    return buildMiniTocFromCollected(collected, 2)
+  }
 }
 
 export default async function renderPage(req: ExtendedRequest, res: Response) {
diff --git a/src/frame/tests/mini-toc-items.ts b/src/frame/tests/mini-toc-items.ts
