Frontend: Fix lazy highlight.js aliases and autodetection

src/Elastic.Documentation.Site/Assets/hljs.test.ts

@@ -7,12 +7,14 @@ async function loadModule(): Promise<HljsModule> {
     return await import('./hljs')
 }
 
-function setCodeBlocks(...languages: string[]) {
+function setCodeBlocks(...languages: (string | null)[]) {
     const blocks = languages
-        .map(
-            (language) =>
-                `<pre><code class="language-${language}">echo hello</code></pre>`
-        )
+        .map((language) => {
+            if (language === null) {
+                return '<pre><code>const x = 1</code></pre>'
+            }
+            return `<pre><code class="language-${language}">echo hello</code></pre>`
+        })
         .join('')
     document.body.innerHTML = `<div id="markdown-content">${blocks}</div>`
 }
@@ -22,14 +24,15 @@ afterEach(() => {
 })
 
 describe('initHighlight', () => {
-    it('loads only the languages referenced on the page', async () => {
+    it('loads all curated languages when any code block is present', async () => {
         const { initHighlight, hljs } = await loadModule()
         setCodeBlocks('bash')
 
         await initHighlight()
 
         expect(hljs.listLanguages()).toContain('bash')
-        expect(hljs.listLanguages()).not.toContain('python')
+        expect(hljs.listLanguages()).toContain('python')
+        expect(hljs.listLanguages()).toContain('javascript')
     })
 
     it('loads no node_modules language when there are no code blocks', async () => {
@@ -52,26 +55,45 @@ describe('initHighlight', () => {
         expect(block?.getAttribute('data-highlighted')).toBe('yes')
     })
 
-    it('loads a newly encountered language on a later swap', async () => {
-        const { initHighlight, hljs } = await loadModule()
-        setCodeBlocks('bash')
-        await initHighlight()
-        expect(hljs.listLanguages()).not.toContain('python')
+    it.each([
+        ['js', 'javascript'],
+        ['ts', 'typescript'],
+        ['jsx', 'javascript'],
+    ])(
+        'highlights alias language-%s via %s aliases',
+        async (alias, canonical) => {
+            const { initHighlight } = await loadModule()
+            setCodeBlocks(alias)
+
+            await initHighlight()
+
+            const block = document.querySelector('#markdown-content pre code')
+            expect(block?.getAttribute('data-highlighted')).toBe('yes')
+            expect(block?.className).toContain(`language-${canonical}`)
+        }
+    )
+
+    it('autodetects an unlabeled code block', async () => {
+        const { initHighlight } = await loadModule()
+        setCodeBlocks(null)
 
-        // Simulate an htmx swap introducing a language not present initially.
-        setCodeBlocks('python')
         await initHighlight()
 
-        expect(hljs.listLanguages()).toContain('python')
+        const block = document.querySelector('#markdown-content pre code')
+        expect(block?.getAttribute('data-highlighted')).toBe('yes')
+        expect(block?.querySelector('.hljs-keyword')).not.toBeNull()
     })
 
-    it('resolves aliases to their canonical language', async () => {
+    it('resolves sh alias to shell', async () => {
         const { initHighlight, hljs } = await loadModule()
         setCodeBlocks('sh')
 
         await initHighlight()
 
         expect(hljs.listLanguages()).toContain('shell')
+        const block = document.querySelector('#markdown-content pre code')
+        expect(block?.getAttribute('data-highlighted')).toBe('yes')
+        expect(block?.className).toContain('language-shell')
     })
 
     it('does not re-highlight a block when invoked concurrently', async () => {
@@ -100,6 +122,28 @@ describe('initHighlight', () => {
     })
 })
 
+describe('highlightCodeBlocks', () => {
+    it('scopes highlighting to the supplied root', async () => {
+        const { highlightCodeBlocks } = await loadModule()
+        document.body.innerHTML = `
+            <div id="inside"><pre><code class="language-bash">echo inside</code></pre></div>
+            <div id="outside"><pre><code class="language-bash">echo outside</code></pre></div>
+        `
+
+        const inside = document.querySelector('#inside')!
+        await highlightCodeBlocks(inside)
+
+        expect(
+            inside.querySelector('code')?.getAttribute('data-highlighted')
+        ).toBe('yes')
+        expect(
+            document
+                .querySelector('#outside code')
+                ?.getAttribute('data-highlighted')
+        ).not.toBe('yes')
+    })
+})
+
 describe('toLanguageFn', () => {
     // Parcel's dynamic import() resolves a language module to the LanguageFn directly,
     // while Babel/Jest wrap it as { default: fn }. Both must yield the function; a

src/Elastic.Documentation.Site/Assets/hljs.ts

@@ -1,7 +1,6 @@
 import { mergeHTMLPlugin } from './hljs-merge-html-plugin'
 import { LanguageFn } from 'highlight.js'
 import hljs from 'highlight.js/lib/core'
-import { $$optional } from 'select-dom'
 
 // highlight.js language modules and the esql plugin default-export the LanguageFn.
 // Parcel's dynamic import() resolves to the module's exports directly (the function),
@@ -13,8 +12,7 @@ export function toLanguageFn(mod: unknown): LanguageFn {
 }
 
 // Each entry lazily imports one highlight.js language module (or the esql plugin) so
-// only the languages actually present on a page are downloaded, instead of eagerly
-// bundling all of them into the entry chunk.
+// grammars stay out of the entry chunk until a page or message contains code.
 const languageLoaders: Record<string, () => Promise<LanguageFn>> = {
     asciidoc: () =>
         import('highlight.js/lib/languages/asciidoc').then(toLanguageFn),
@@ -67,37 +65,39 @@ const languageLoaders: Record<string, () => Promise<LanguageFn>> = {
     yaml: () => import('highlight.js/lib/languages/yaml').then(toLanguageFn),
 }
 
-// Alias -> canonical language name. Aliases are registered with hljs once their
-// canonical module has loaded so `language-sh` etc. resolve correctly.
-const languageAliases: Record<string, string> = {
-    sh: 'shell',
+const CODE_BLOCK_SELECTOR = 'pre code:not([data-highlighted])'
+
+let allLanguagesReady: Promise<void> | null = null
+
+function ensureHighlightReady(): Promise<void> {
+    if (allLanguagesReady) return allLanguagesReady
+
+    allLanguagesReady = Promise.all(
+        Object.entries(languageLoaders).map(async ([name, loader]) => {
+            hljs.registerLanguage(name, toLanguageFn(await loader()))
+        })
+    ).then(() => {
+        hljs.registerAliases(['sh'], { languageName: 'shell' })
+    })
+
+    return allLanguagesReady
 }
 
-// Caches the registration promise per canonical language so concurrent code blocks
-// (and later htmx swaps) share a single import instead of re-fetching the module.
-const registrations = new Map<string, Promise<void>>()
+export async function highlightCodeBlocks(root: ParentNode): Promise<void> {
+    const blocks = root.querySelectorAll<HTMLElement>(CODE_BLOCK_SELECTOR)
+    if (blocks.length === 0) return
 
-// Imports and registers a language (plus any aliases pointing at it) exactly once.
-// Unknown languages have no loader and resolve immediately so callers can await
-// unconditionally.
-function ensureLanguage(name: string): Promise<void> {
-    const canonical = languageAliases[name] ?? name
-    const loader = languageLoaders[canonical]
-    if (!loader) return Promise.resolve()
+    await ensureHighlightReady()
 
-    const cached = registrations.get(canonical)
-    if (cached) return cached
+    for (const block of blocks) {
+        if (!block.dataset.highlighted) hljs.highlightElement(block)
+    }
+}
 
-    const registration = loader().then((languageFn) => {
-        hljs.registerLanguage(canonical, languageFn)
-        for (const [alias, target] of Object.entries(languageAliases)) {
-            if (target === canonical) {
-                hljs.registerAliases([alias], { languageName: canonical })
-            }
-        }
-    })
-    registrations.set(canonical, registration)
-    return registration
+export async function initHighlight(): Promise<void> {
+    const root = document.querySelector('#markdown-content')
+    if (!root) return
+    await highlightCodeBlocks(root)
 }
 
 hljs.registerLanguage('apiheader', function () {
@@ -230,39 +230,5 @@ hljs.addPlugin(mergeHTMLPlugin)
 // for code callouts
 hljs.configure({ ignoreUnescapedHTML: true })
 
-function getLanguageFromClassList(element: Element): string | undefined {
-    for (const className of element.classList) {
-        if (className.startsWith('language-')) {
-            return className.slice('language-'.length)
-        }
-    }
-    return undefined
-}
-
-export async function initHighlight() {
-    const blocks = $$optional(
-        '#markdown-content pre code:not([data-highlighted])'
-    )
-    if (blocks.length === 0) return
-
-    // Import only the language modules referenced by the unprocessed blocks before
-    // highlighting. Unknown/plain-text languages have no loader and resolve immediately;
-    // hljs.highlightElement then degrades gracefully without breaking other blocks.
-    const requiredLanguages = new Set<string>()
-    for (const block of blocks) {
-        const language = getLanguageFromClassList(block)
-        if (language) requiredLanguages.add(language)
-    }
-
-    await Promise.all([...requiredLanguages].map(ensureLanguage))
-
-    // Re-check data-highlighted after awaiting: a concurrent initHighlight (e.g. a
-    // second htmx:load fired while the language import was pending) may already have
-    // highlighted these captured blocks, and hljs warns if asked to redo one.
-    blocks.forEach((block) => {
-        if (!block.dataset.highlighted) hljs.highlightElement(block)
-    })
-}
-
 // Export the configured hljs instance for reuse
 export { hljs }

src/Elastic.Documentation.Site/Assets/web-components/AskAi/ChatMessage.tsx

@@ -1,7 +1,8 @@
 import { initCopyButton } from '../../copybutton'
-import { hljs } from '../../hljs'
+import { highlightCodeBlocks } from '../../hljs'
 import { ApiError } from '../shared/errorHandling'
 import { useHtmxContainer } from '../shared/htmx/useHtmxContainer'
+import { markdownRenderer } from '../shared/markdownRenderer'
 import { AskAiEvent, ChunkEvent, EventTypes } from './AskAiEvent'
 import { aiGradients } from './ElasticAiAssistantButton'
 import { GeneratingStatus } from './GeneratingStatus'
@@ -23,39 +24,8 @@ import {
 } from '@elastic/eui'
 import { css } from '@emotion/react'
 import DOMPurify from 'dompurify'
-import { Marked, RendererObject, Tokens } from 'marked'
 import { useEffect, useMemo, useRef } from 'react'
 
-// Create the marked instance once globally (renderer never changes)
-const createMarkedInstance = () => {
-    const renderer: RendererObject = {
-        code({ text, lang }: Tokens.Code): string {
-            let highlighted: string
-            try {
-                highlighted = lang
-                    ? hljs.highlight(text, { language: lang }).value
-                    : hljs.highlightAuto(text).value
-            } catch {
-                // Fallback to auto highlighting if the specified language is not found
-                highlighted = hljs.highlightAuto(text).value
-            }
-            return `<div class="highlight">
-                <pre>
-                    <code class="language-${lang}">${highlighted}</code>
-                </pre>
-            </div>`
-        },
-        table(token: Tokens.Table): string {
-            const defaultMarked = new Marked()
-            const defaultTableHtml = defaultMarked.parse(token.raw)
-            return `<div class="table-wrapper">${defaultTableHtml}</div>`
-        },
-    }
-    return new Marked({ renderer })
-}
-
-const markedInstance = createMarkedInstance()
-
 interface ChatMessageProps {
     message: ChatMessageType
     events?: AskAiEvent[]
@@ -387,7 +357,7 @@ export const ChatMessage = ({
     }, [content, isComplete])
 
     const parsed = useMemo(() => {
-        const html = markedInstance.parse(mainContent) as string
+        const html = markdownRenderer.parse(mainContent, { async: false })
         return DOMPurify.sanitize(html)
     }, [mainContent])
 
@@ -402,21 +372,26 @@ export const ChatMessage = ({
     const ref = useRef<HTMLDivElement>(null)
 
     useEffect(() => {
-        if (isComplete && ref.current) {
-            const timer = setTimeout(() => {
-                try {
-                    initCopyButton(
-                        '.highlight pre',
-                        ref.current!,
-                        'ai-message-codecell-'
-                    )
-                } catch (error) {
-                    console.error('Failed to initialize copy buttons:', error)
-                }
-            }, 100)
-            return () => clearTimeout(timer)
+        if (!ref.current || !parsed) return
+
+        let cancelled = false
+        void highlightCodeBlocks(ref.current).then(() => {
+            if (cancelled || !isComplete || !ref.current) return
+            try {
+                initCopyButton(
+                    '.highlight pre',
+                    ref.current,
+                    'ai-message-codecell-'
+                )
+            } catch (error) {
+                console.error('Failed to initialize copy buttons:', error)
+            }
+        })
+
+        return () => {
+            cancelled = true
         }
-    }, [isComplete])
+    }, [parsed, isComplete])
 
     // Process internal docs links for htmx navigation
     useHtmxContainer(ref, [parsed])

src/Elastic.Documentation.Site/Assets/web-components/shared/MarkdownContent.test.tsx

@@ -0,0 +1,40 @@
+import { MarkdownContent } from './MarkdownContent'
+import { EuiProvider } from '@elastic/eui'
+import { render, waitFor } from '@testing-library/react'
+
+jest.mock('../../copybutton', () => ({
+    initCopyButton: jest.fn(),
+}))
+
+const renderMarkdown = (content: string) =>
+    render(
+        <EuiProvider
+            colorMode="light"
+            globalStyles={false}
+            utilityClasses={false}
+        >
+            <MarkdownContent content={content} enableCopyButtons={false} />
+        </EuiProvider>
+    )
+
+describe('MarkdownContent', () => {
+    it('highlights alias fences after mount', async () => {
+        renderMarkdown('```js\nconst x = 1\n```')
+
+        await waitFor(() => {
+            const block = document.querySelector('.markdown-content pre code')
+            expect(block?.getAttribute('data-highlighted')).toBe('yes')
+            expect(block?.className).toContain('language-javascript')
+        })
+    })
+
+    it('autodetects unlabeled fences after mount', async () => {
+        renderMarkdown('```\nconst x = 1\n```')
+
+        await waitFor(() => {
+            const block = document.querySelector('.markdown-content pre code')
+            expect(block?.getAttribute('data-highlighted')).toBe('yes')
+            expect(block?.querySelector('.hljs-keyword')).not.toBeNull()
+        })
+    })
+})