fix(app): a math or latex code block of equations draws as typeset math

app/e2e/math.spec.ts

@@ -46,3 +46,39 @@ test("a reply's math draws as typeset math, and a single $ stays text", async ({
   // A formula KaTeX can't read is its source, and the rest still draws.
   await expect(md.locator(".katex-error")).toHaveText("\\frac{1}{");
 });
+
+const BLOCKS_REPLY = [
+  "## Rigid bodies, step by step",
+  "",
+  "Semi-implicit Euler, velocity first:",
+  "",
+  "```math",
+  "\\begin{aligned} \\mathbf{v}_{t+\\Delta t} &= \\mathbf{v}_t + \\Delta t \\, M^{-1} \\mathbf{F} \\\\ \\mathbf{x}_{t+\\Delta t} &= \\mathbf{x}_t + \\Delta t \\, \\mathbf{v}_{t+\\Delta t} \\end{aligned}",
+  "```",
+  "",
+  "The contact impulse, clamped so bodies only push:",
+  "",
+  "```latex",
+  "\\lambda = \\max\\left(0, -\\frac{J\\mathbf{v} + b}{J M^{-1} J^{\\top}}\\right)",
+  "```",
+  "",
+  "And the file to compile, which stays code:",
+  "",
+  "```latex",
+  "\\documentclass{article}",
+  "\\begin{document}",
+  "Rigid bodies.",
+  "\\end{document}",
+  "```",
+].join("\n");
+
+test("a ```math block and a ```latex one of equations draw; a LaTeX document stays code", async ({ app }) => {
+  mockOnly("lands a reply on the demo's agent");
+  await app.open({ params: { view: "conversation" } });
+  await app.openWorktree("gpu/judge-v2");
+  await landReply(app.page, BLOCKS_REPLY);
+  const md = app.chat.locator("[data-kind=text] .cv-md").filter({ hasText: "Rigid bodies, step by step" });
+  await expect(md.locator(".cv-math-block .katex-display")).toHaveCount(2);
+  await expect(md.locator(".katex-error")).toHaveCount(0);
+  await expect(md.locator("pre code", { hasText: "\\documentclass{article}" })).toBeVisible();
+});

app/src/lib/markdown-html.test.ts

@@ -53,6 +53,32 @@ test("math in code stays as written", () => {
   assert.equal(read("[docs](https://example.com/a\\(b\\))"), '<p><a href="https://example.com/a(b)">docs</a></p>');
 });
 
+test("a ```math block, or a ```latex one of equations, is a display", () => {
+  const block = (tex: string) => `<div class="cv-math-block"><span class="cv-math cv-math-display">[display:${tex}]</span></div>`;
+  assert.equal(read("```math\nE = mc^2\n```"), block("E = mc^2"));
+  assert.equal(read("```latex\n\\begin{aligned}a &= b \\\\ c &= d\\end{aligned}\n```"), block("\\begin{aligned}a &= b \\\\ c &= d\\end{aligned}"));
+  assert.equal(read("```tex\nx^2\n```"), block("x^2"));
+  // What micromark escaped is TeX again.
+  assert.equal(read("```math\na < b & c\n```"), block("a < b & c"));
+  assert.match(html("```math\n\\frac{1}{2}\n```"), /katex-display/);
+  // Words around it are still the reply's.
+  assert.equal(read("Energy:\n\n```math\nE = mc^2\n```\n\nas Einstein had it."), `<p>Energy:</p>\n${block("E = mc^2")}\n<p>as Einstein had it.</p>`);
+});
+
+test("a LaTeX document, other code and a block still being written stay code", () => {
+  // A document is for pdflatex.
+  const doc = "```latex\n\\documentclass{article}\n\\begin{document}\n$x$\n\\end{document}\n```";
+  assert.equal(markdownHtml(doc, tex).math, false);
+  assert.match(read(doc), /<pre><code class="language-latex">\\documentclass/);
+  assert.match(read("```tex\n\\section{Forces}\nF = ma\n```"), /<pre><code class="language-tex">/);
+  // Other languages and unnamed blocks are untouched.
+  assert.equal(read("```python\nx = 1\n```"), '<pre><code class="language-python">x = 1\n</code></pre>');
+  // A draft's block that hasn't closed: code until it does.
+  assert.equal(markdownHtml("```math\nE = m", tex).math, false);
+  // Only the open one, when one before it closed.
+  assert.equal(read("```math\na\n```\n\n```math\nb"), '<div class="cv-math-block"><span class="cv-math cv-math-display">[display:a]</span></div>\n<pre><code class="language-math">b\n</code></pre>\n');
+});
+
 test("a formula KaTeX can't read shows as its source, and nothing throws", () => {
   const out = html("A \\(\\frac{1}{\\) and $$\\badcommand$$ here.");
   assert.match(out, /katex-error/);

app/src/lib/markdown-html.ts

@@ -106,6 +106,16 @@ const MARK = (n: number) => `﷐${n}﷑`;
 const MARKS = /﷐(\d+)﷑/g;
 
 const escapeHtml = (s: string) => s.replace(/&/g, "&amp;").replace(/</g, "&lt;").replace(/>/g, "&gt;").replace(/"/g, "&quot;");
+const unescapeHtml = (s: string) => s.replace(/&lt;/g, "<").replace(/&gt;/g, ">").replace(/&quot;/g, '"').replace(/&amp;/g, "&");
+
+// A code block of math: ```math, as GitHub writes a display, or ```latex
+// and ```tex holding only equations. One that is a document (to compile
+// with pdflatex, not to read here) stays code, as does one still open in
+// a draft: a half-written formula would only show KaTeX's errors.
+const MATH_LANGS = "math|latex|tex|katex";
+const MATH_CODE = new RegExp(`<pre><code class="language-(${MATH_LANGS})">([\\s\\S]*?)<\\/code><\\/pre>`, "gi");
+const MATH_FENCE_CLOSED = new RegExp(`^ {0,3}(\`{3,}|~{3,})[ \\t]*(?:${MATH_LANGS})\\b[^\\n]*\\n[\\s\\S]*?\\n {0,3}\\1[ \\t]*$`, "gim");
+const TEX_DOCUMENT = /\\(documentclass|usepackage|begin\{document\}|section|subsection|title|maketitle|author|chapter|item)\b/;
 
 interface Formula {
   tex: string;
@@ -194,6 +204,14 @@ export function markdownHtml(text: string, tex?: TexRender, shape: (html: string
     if (!unclosed.length || unclosed[0] < 0) break;
     source = `${source.slice(0, unclosed[0])}\\${source.slice(unclosed[0])}`;
   }
+  let closed = (source.match(MATH_FENCE_CLOSED) ?? []).length;
+  html = html.replace(MATH_CODE, (all, lang: string, body: string) => {
+    if (closed-- <= 0) return all;
+    const src = unescapeHtml(body).replace(/\n$/, "");
+    if (!src.trim() || (lang.toLowerCase() !== "math" && TEX_DOCUMENT.test(src))) return all;
+    formulas.push({ tex: src, display: true, source: `\`\`\`${lang}\n${src}\n\`\`\`` });
+    return `<div class="cv-math-block">${MARK(formulas.length - 1)}</div>`;
+  });
   if (!formulas.length) return { html: shape(html), math: false };
   // $$…$$ on a line of its own is a display, as a $$ block is.
   html = html.replace(/<p>﷐(\d+)﷑<\/p>/g, (all, n) => {

docs/changelog.mdx

@@ -11,6 +11,10 @@ again.
 
 ## Unreleased
 
+- **Math in code blocks draws too.** A ` ```math ` block, as GitHub writes a
+  formula, and a ` ```latex ` or ` ```tex ` block of equations now show as
+  typeset math. A whole LaTeX document (with `\documentclass`, meant for
+  `pdflatex`) stays code, as does a block still being written.
 - **Agents started by other agents follow Structured agents too.** With
   *Structured agents* on in Settings › Labs, an agent that another agent
   starts with the berth tools, one started with `berth task new` or