refactor(desktop): move markdown rendering to rust (#10000)

1 files changed, 101 insertions, 2 deletions

diff --git a/packages/ui/src/context/marked.tsx b/packages/ui/src/context/marked.tsx
index 6cf1dd54e..71881353a 100644
--- a/packages/ui/src/context/marked.tsx
+++ b/packages/ui/src/context/marked.tsx
@@ -1,6 +1,7 @@
 import { marked } from "marked"
 import markedKatex from "marked-katex-extension"
 import markedShiki from "marked-shiki"
+import katex from "katex"
 import { bundledLanguages, type BundledLanguage } from "shiki"
 import { createSimpleContext } from "./helper"
 import { getSharedHighlighter, registerCustomTheme, ThemeRegistrationResolved } from "@pierre/diffs"
@@ -375,10 +376,95 @@ registerCustomTheme("OpenCode", () => {
   } as unknown as ThemeRegistrationResolved)
 })
 
+function renderMathInText(text: string): string {
+  let result = text
+
+  // Display math: $$...$$
+  const displayMathRegex = /\$\$([\s\S]*?)\$\$/g
+  result = result.replace(displayMathRegex, (_, math) => {
+    try {
+      return katex.renderToString(math, {
+        displayMode: true,
+        throwOnError: false,
+      })
+    } catch {
+      return `$$${math}$$`
+    }
+  })
+
+  // Inline math: $...$
+  const inlineMathRegex = /(?<!\$)\$(?!\$)((?:[^$\\]|\\.)+?)\$(?!\$)/g
+  result = result.replace(inlineMathRegex, (_, math) => {
+    try {
+      return katex.renderToString(math, {
+        displayMode: false,
+        throwOnError: false,
+      })
+    } catch {
+      return `$${math}$`
+    }
+  })
+
+  return result
+}
+
+function renderMathExpressions(html: string): string {
+  // Split on code/pre/kbd tags to avoid processing their contents
+  const codeBlockPattern = /(<(?:pre|code|kbd)[^>]*>[\s\S]*?<\/(?:pre|code|kbd)>)/gi
+  const parts = html.split(codeBlockPattern)
+
+  return parts
+    .map((part, i) => {
+      // Odd indices are the captured code blocks - leave them alone
+      if (i % 2 === 1) return part
+      // Process math only in non-code parts
+      return renderMathInText(part)
+    })
+    .join("")
+}
+
+async function highlightCodeBlocks(html: string): Promise<string> {
+  const codeBlockRegex = /<pre><code(?:\s+class="language-([^"]*)")?>([\s\S]*?)<\/code><\/pre>/g
+  const matches = [...html.matchAll(codeBlockRegex)]
+  if (matches.length === 0) return html
+
+  const highlighter = await getSharedHighlighter({ themes: ["OpenCode"], langs: [] })
+
+  let result = html
+  for (const match of matches) {
+    const [fullMatch, lang, escapedCode] = match
+    const code = escapedCode
+      .replace(/&lt;/g, "<")
+      .replace(/&gt;/g, ">")
+      .replace(/&amp;/g, "&")
+      .replace(/&quot;/g, '"')
+      .replace(/&#39;/g, "'")
+
+    let language = lang || "text"
+    if (!(language in bundledLanguages)) {
+      language = "text"
+    }
+    if (!highlighter.getLoadedLanguages().includes(language)) {
+      await highlighter.loadLanguage(language as BundledLanguage)
+    }
+
+    const highlighted = highlighter.codeToHtml(code, {
+      lang: language,
+      theme: "OpenCode",
+      tabindex: false,
+    })
+    result = result.replace(fullMatch, () => highlighted)
+  }
+
+  return result
+}
+
+export type NativeMarkdownParser = (markdown: string) => Promise<string>
+
 export const { use: useMarked, provider: MarkedProvider } = createSimpleContext({
   name: "Marked",
-  init: () => {
-    return marked.use(
+  init: (props: { nativeParser?: NativeMarkdownParser }) => {
+    const jsParser = marked.use(
       {
         renderer: {
           link({ href, title, text }) {
@@ -407,5 +493,18 @@ export const { use: useMarked, provider: MarkedProvider } = createSimpleContext(
         },
       }),
     )
+
+    if (props.nativeParser) {
+      const nativeParser = props.nativeParser
+      return {
+        async parse(markdown: string): Promise<string> {
+          const html = await nativeParser(markdown)
+          const withMath = renderMathExpressions(html)
+          return highlightCodeBlocks(withMath)
+        },
+      }
+    }
+
+    return jsParser
   },
 })