fix(ui): render dollar math with marked-katex-extension

packages/session-ui/src/components/markdown-stream.test.ts

@@ -18,6 +18,13 @@ describe("markdown stream", () => {
     expect(stream(raw, true)).toEqual([{ raw, src, mode: "live" }])
   })
 
+  test("keeps a display math block with a blank line in one stable block", () => {
+    expect(stream("$$\na\n\nb\n$$\n\nafter", true)).toEqual([
+      { raw: "$$\na\n\nb\n$$\n\n", src: "$$\na\n\nb\n$$\n\n", mode: "full" },
+      { raw: "after", src: "after", mode: "live" },
+    ])
+  })
+
   test("keeps incomplete links non-clickable until they finish", () => {
     expect(stream("see [docs](https://example.com/gu", true)).toEqual([
       { raw: "see [docs](https://example.com/gu", src: "see docs", mode: "live" },

packages/session-ui/src/components/markdown-stream.ts

@@ -1,5 +1,6 @@
-import { marked, type Tokens } from "marked"
+import { Marked, type Tokens } from "marked"
 import remend from "remend"
+import { markdownMath } from "@opencode/ui/context/marked-parser"
 import { completedProjection } from "./markdown-projection"
 
 export type Block = {
@@ -15,6 +16,9 @@ export type Projection = {
   blocks: Block[]
 }
 
+// Split with the renderer's math syntax so a display block is never cut into separately rendered pieces.
+const lexer = new Marked(...markdownMath)
+
 function refs(text: string) {
   if (!text.includes("]:")) return false
 
@@ -61,7 +65,7 @@ export function stream(text: string, live: boolean): Block[] {
   if (!live) return completedProjection(text).blocks
 
   if (refs(text)) return [{ raw: text, src: heal(text), mode: "live" }] satisfies Block[]
-  const tokens = marked.lexer(text)
+  const tokens = lexer.lexer(text)
   const tail = tokens.findLastIndex((token) => token.type !== "space")
 
   if (tail < 0) return [{ raw: text, src: heal(text), mode: "live" }] satisfies Block[]
@@ -80,6 +84,7 @@ export function stream(text: string, live: boolean): Block[] {
     while (tokens[index + 1]?.type === "space" && index + 1 < tail) raw += tokens[++index]!.raw
 
     if (token.type === "code") {
+      // SAFETY: only marked's built-in fence tokenizer emits `code`; the math extensions use their own token types.
       const code = token as Tokens.Code
       result.push({ raw, src: code.text, mode: "code", language: language(code.lang), complete: true })
       continue
@@ -95,6 +100,7 @@ export function stream(text: string, live: boolean): Block[] {
 
   if (last.type !== "code") return [...result, { raw, src: heal(raw), mode: "live" }]
 
+  // SAFETY: only marked's built-in fence tokenizer emits `code`; the math extensions use their own token types.
   const code = last as Tokens.Code
 
   if (!open(code.raw))

packages/ui/package.json

@@ -116,6 +116,7 @@
     "fuzzysort": "catalog:",
     "katex": "0.16.47",
     "marked": "catalog:",
+    "marked-katex-extension": "5.1.13",
     "marked-shiki": "catalog:",
     "motion": "12.34.5",
     "remeda": "catalog:",

packages/ui/src/context/marked-base.tsx

@@ -16,7 +16,8 @@ let smallParser: Marked | undefined
 
 export function parseSmallMarkdown(text: string) {
   // Any possible KaTeX delimiter stays on the worker, including escaped ones.
-  if (text.length > 1024 || text.includes("\\(") || text.includes("$$")) return
+  // Dollar math needs a pair, so a lone price such as `$5` can stay on this path.
+  if (text.length > 1024 || text.includes("\\(") || text.indexOf("$") !== text.lastIndexOf("$")) return
 
   // Ordinary prose does not need the block/inline lexer's cold regular expressions.
   // Keep every possible Markdown construct, autolink, and hard break on that lexer.

packages/ui/src/context/marked-parser.test.ts

@@ -15,6 +15,42 @@ test("renders inline and block math", async () => {
   expect(await parser.parse("$$\nx^2\n$$\n")).toContain('<span class="katex-display">')
 })
 
+test.each(["Energy is $E = mc^2$.", "因此 $x^2$ 是正数", "$a_1$ and $b_1$."])(
+  "renders dollar inline math: %s",
+  async (text) => {
+    const html = await parser.parse(text)
+    expect(html).toContain('<span class="katex">')
+    expect(html).not.toContain("$")
+  },
+)
+
+test.each(["$$E = mc^2$$", "Inline $$x$$ display", "$$\na\n\nb\n$$\n"])(
+  "renders dollar display math: %s",
+  async (text) => {
+    const html = await parser.parse(text)
+    expect(html).toContain('<span class="katex-display">')
+    expect(html).not.toContain("$$")
+  },
+)
+
+test.each([
+  "It costs $5 and $10.",
+  "From $5-$10 a month",
+  "It costs $5, originally $10, ranging from $1,000 to $2,000.",
+  "Use $HOME and $PATH",
+  "Use `$x$` as a placeholder",
+])("leaves prices, shell variables, and code as text: %s", async (text) => {
+  expect(await parser.parse(text)).not.toContain("katex")
+})
+
+// marked-katex-extension requires a space or line start before `$` and a space or punctuation after the closing `$`.
+test.each(["where (i.e. $x$) holds", "the $x$-axis", "因此$x^2$是正数"])(
+  "leaves tight dollar math as text: %s",
+  async (text) => {
+    expect(await parser.parse(text)).not.toContain("katex")
+  },
+)
+
 test("uses the configured code highlighter", async () => {
   expect(await parser.parse("```ts\nconst value = 1\n```\n")).toBe('<pre data-language="ts">const value = 1</pre>\n')
 })
@@ -51,6 +87,7 @@ test.each([
   "foo\n```",
   "\\(x^2\\)",
   "$$\nx^2\n$$\n",
+  "Energy is $E = mc^2$.",
   "a".repeat(1025),
 ])("leaves code, math, and large Markdown to the worker: %s", (text) => {
   expect(parseSmallMarkdown(text)).toBeUndefined()

packages/ui/src/context/marked-parser.tsx

@@ -1,65 +1,45 @@
 import katex from "katex"
 import type { MarkedExtension, Tokens } from "marked"
+import markedKatex from "marked-katex-extension"
 import markedShiki from "marked-shiki"
 import { createMarkdownBase } from "./marked-base"
 
 export function createMarkdownParser(highlight: (code: string, language: string) => string | Promise<string>) {
-  return createMarkdownBase().use(katexExtension, markedShiki({ highlight }))
+  return createMarkdownBase().use(...markdownMath, markedShiki({ highlight }))
 }
 
-const inlineMathRegex = /^\\\(((?:\\.|[^\\\n])*?)\\\)/
-
-const blockMathRegex = /^\$\$\n([\s\S]+?)\n\$\$(?:\n|$)/
-
-const katexExtension: MarkedExtension = {
-  extensions: [
-    {
-      name: "inlineKatex",
-      level: "inline",
-      start(src) {
-        const index = src.indexOf("\\(")
-
-        if (index === -1) return
-
-        return index
-      },
-      tokenizer(src) {
-        const match = src.match(inlineMathRegex)
-
-        if (!match) return
-
-        return {
-          type: "inlineKatex",
-          raw: match[0],
-          text: match[1].trim(),
-          displayMode: false,
-        }
+const inlineParenMathRegex = /^\\\(((?:\\.|[^\\\n])*?)\\\)/
+
+// marked-katex-extension handles `$...$`, `$$...$$`, and `$$` fenced blocks; it has no `\(...\)` syntax.
+export const markdownMath: MarkedExtension[] = [
+  markedKatex({ throwOnError: false }),
+  {
+    extensions: [
+      {
+        name: "inlineParenKatex",
+        level: "inline",
+        start(src) {
+          const index = src.indexOf("\\(")
+
+          if (index === -1) return
+
+          return index
+        },
+        tokenizer(src) {
+          const match = src.match(inlineParenMathRegex)
+
+          if (!match) return
+
+          return {
+            type: "inlineParenKatex",
+            raw: match[0],
+            text: match[1].trim(),
+          }
+        },
+        renderer(token: Tokens.Generic) {
+          return katex.renderToString(token.text, { throwOnError: false })
+        },
       },
-      renderer: renderKatexToken,
-    },
-    {
-      name: "blockKatex",
-      level: "block",
-      tokenizer(src) {
-        const match = src.match(blockMathRegex)
-
-        if (!match) return
-
-        return {
-          type: "blockKatex",
-          raw: match[0],
-          text: match[1].trim(),
-          displayMode: true,
-        }
-      },
-      renderer: renderKatexToken,
-    },
-  ],
-}
-
-function renderKatexToken(token: Tokens.Generic) {
-  return katex.renderToString(typeof token.text === "string" ? token.text : "", {
-    displayMode: token.displayMode === true,
-    throwOnError: false,
-  })
-}
+    ],
+  },
+]