1import type { Nodes } from "mdast";
2import { describe, expect, it } from "vite-plus/test";
3
4import { mathKind } from "./markdownMath.ts";
5import { markdownMathRanges, parseMarkdownWithMath } from "./markdownMath.ts";
6
7const parse = parseMarkdownWithMath;
8
9function collect(node: Nodes, out: Nodes[] = []): Nodes[] {
10 out.push(node);
11 if ("children" in node) for (const child of node.children) collect(child, out);
12 return out;
13}
14
15/** Formulas as `[kind, tex]`, in document order. */
16function formulas(markdown: string) {
17 return collect(parse(markdown)).flatMap((node) => {
18 if (node.type === "math") return [["display", node.value]];
19 if (node.type !== "inlineMath") return [];
20 const className = node.data?.hProperties?.className;
21 return [[mathKind(Array.isArray(className) ? className.join(" ") : undefined), node.value]];
22 });
23}
24
25function text(markdown: string) {
26 return collect(parse(markdown))
27 .flatMap((node) => (node.type === "text" ? [node.value] : []))
28 .join("");
29}
30
31describe("chat math parsing", () => {
32 it("keeps GFM autolink destinations intact while restoring literal dollars", () => {
33 const source = "Query www.example.com/?$select=id and $x$.";
34 const link = collect(parse(source)).find((node) => node.type === "link");
35 expect(link).toMatchObject({ url: "http://www.example.com/?$select=id" });
36 expect(formulas(source)).toEqual([["inline", "x"]]);
37 });
38
39 it.each(["", "", ""])(
40 "preserves decoded private-use reference %s",
41 (reference) => {
42 const source = `${reference} [${reference}](https://example.com/?q=${reference}) $x$`;
43 const nodes = collect(parse(source));
44 expect(
45 nodes
46 .filter((node) => node.type === "text")
47 .map((node) => node.value)
48 .join(""),
49 ).toBe("\uE000 \uE000 ");
50 expect(nodes.find((node) => node.type === "link")).toMatchObject({
51 url: "https://example.com/?q=\uE000",
52 });
53 expect(formulas(source)).toEqual([["inline", "x"]]);
54 },
55 );
56
57 it.each(["\\(x\\)", "$x$", "$$x$$"])("preserves indented code containing %s", (formula) => {
58 const source = ` ${formula}`;
59 expect(collect(parse(source)).find((node) => node.type === "code")).toMatchObject({
60 value: formula,
61 });
62 expect(markdownMathRanges(source)).toEqual([]);
63 });
64
65 it("parses formulas without reserving private-use characters", () => {
66 const characters = Array.from({ length: 6400 }, (_, index) =>
67 String.fromCharCode(0xe000 + index),
68 ).join("");
69 const source = "$x$ " + characters;
70 expect(formulas(source)).toEqual([["inline", "x"]]);
71 expect(text(source)).toBe(" " + characters);
72 });
73
74 it("only exposes closed explicit math fences to editors and native rendering", () => {
75 expect(markdownMathRanges("```math\nx^2")).toEqual([]);
76 expect(markdownMathRanges("```math\nx^2\n```")[0]?.tex).toBe("x^2");
77 expect(markdownMathRanges("> ```math\n> x^2\n> ```")[0]?.tex).toBe("x^2");
78 });
79 it("reads the four delimiters models emit", () => {
80 expect(formulas("Inline $x^2$ and \\(y_1\\), display $$a+b$$ and \\[c\\].")).toEqual([
81 ["inline", "x^2"],
82 ["inline", "y_1"],
83 ["display", "a+b"],
84 ["display", "c"],
85 ]);
86 });
87
88 it("reads display blocks on their own lines", () => {
89 expect(formulas("$$\n\\frac{a}{b}\n$$")).toEqual([["display", "\\frac{a}{b}"]]);
90 expect(
91 formulas("Where\n\\[\n\\begin{aligned} a &= b \\\\ c &= d \\end{aligned}\n\\]\nholds."),
92 ).toEqual([["display", "\\begin{aligned} a &= b \\\\ c &= d \\end{aligned}"]]);
93 });
94
95 it.each([" is the first formula.", ", followed by prose.", " \tthen more text."])(
96 "keeps text after a same-line formula separate from later display math: %s",
97 (trailing) => {
98 const source = `$$a$$${trailing}\nSome prose.\n$$\nb\n$$`;
99 const nodes = collect(parse(source));
100 expect(
101 nodes.filter((node) => node.type === "inlineMath" || node.type === "math"),
102 ).toMatchObject([
103 { type: "inlineMath", value: "a" },
104 { type: "math", value: "b" },
105 ]);
106 expect(text(source)).toContain(`${trailing}\nSome prose.`);
107 expect(markdownMathRanges(source).map((formula) => formula.source)).toEqual([
108 "$$a$$",
109 "$$\nb\n$$",
110 ]);
111 },
112 );
113
114 it("keeps block syntax inside display blocks as TeX", () => {
115 expect(formulas("\\[\nx\n- y\n# z\n\\]\nafter")).toEqual([["display", "x\n- y\n# z"]]);
116 expect(formulas("> $$\n> x\n> $$")).toEqual([["display", "x"]]);
117 // Same-line content is not a fence; it stays one formula.
118 expect(formulas("$$x^2\n+y^2$$")).toEqual([["display", "x^2\n+y^2"]]);
119 expect(text("$$E=mc^2")).toBe("$$E=mc^2");
120 });
121
122 it("keeps prices, shell variables, and skill tokens as prose", () => {
123 for (const prose of [
124 "It costs $5 and $10.",
125 "Budget $5-$10 per seat, or $20,000-$30,000 a year.",
126 "Plans are $5/$10 monthly.",
127 "Pay $5,$10 now.",
128 "Run $browser then $deploy on it.",
129 "Set $HOME before running.",
130 "Set $HOME,$PATH before running.",
131 "Set $HOME, $PATH before running.",
132 "Set $HOME,$_PATH before running.",
133 "Set $home,$path before running.",
134 "Set $HOME,${PATH} before running.",
135 "Set $HOME,$? before running.",
136 "Set $HOME,$@ before running.",
137 "Set $HOME,$(pwd) before running.",
138 "Pay $1,000,$2,000 now.",
139 ]) {
140 expect(formulas(prose), prose).toEqual([]);
141 expect(text(prose), prose).toBe(prose);
142 }
143 });
144
145 it("keeps prices as prose when math follows on the same line", () => {
146 for (const price of ["$5-$10", "$5 and $10", "$5/$10", "$20,000-$30,000"]) {
147 const source = `Price ${price} stays literal. The formula is $x_i$ and $5$.`;
148 expect(formulas(source)).toEqual([
149 ["inline", "x_i"],
150 ["inline", "5"],
151 ]);
152 expect(text(source)).toBe(`Price ${price} stays literal. The formula is and .`);
153 }
154 });
155
156 it("reads formulas that end in a comma", () => {
157 for (const boundary of ["", " ", "\n", ".", ",", ";", ":", ")", "]", "}"]) {
158 expect(formulas(`$x,$${boundary}`)).toEqual([["inline", "x,"]]);
159 }
160 expect(formulas("Where $x = 1,$ with $$a,$$ and $$ b,$$.")).toEqual([
161 ["inline", "x = 1,"],
162 ["display", "a,"],
163 ["display", " b,"],
164 ]);
165 });
166
167 it("leaves code, escapes, links, and citations alone", () => {
168 for (const [markdown, rendered] of [
169 ["Use `$x$` in code.", "Use in code."],
170 ["It costs $5 and $10, and `$HOME` stays code.", "It costs $5 and $10, and stays code."],
171 ["Export $PATH, then `$HOME`.", "Export $PATH, then ."],
172 ["\\(a `\\)` b", "(a b"],
173 [
174 "Set $TOKEN, then open [users](https://graph.microsoft.com/v1.0/users?$select=id).",
175 "Set $TOKEN, then open users.",
176 ],
177 [
178 "Set $TOKEN, then open <https://graph.microsoft.com/v1.0/users?$select=id>.",
179 "Set $TOKEN, then open https://graph.microsoft.com/v1.0/users?$select=id.",
180 ],
181 ["```\n$x$ \\(y\\)\n```", ""],
182 ["Escaped \\$x\\$ stays.", "Escaped $x$ stays."],
183 ["Escaped \\\\(x\\\\) stays.", "Escaped \\(x\\) stays."],
184 ["See [docs](https://example.com/\\(a\\)) here.", "See docs here."],
185 ["As shown in \\[1\\] and \\[2, 3\\].", "As shown in [1] and [2, 3]."],
186 ] as const) {
187 expect(formulas(markdown), markdown).toEqual([]);
188 expect(text(markdown), markdown).toBe(rendered);
189 }
190 });
191
192 it("starts a formula right after an escaped dollar", () => {
193 expect(formulas("Pay \\$$x$ now.")).toEqual([["inline", "x"]]);
194 });
195
196 it("starts a formula right after an escaped backslash", () => {
197 const markdown = "Path \\\\\\(x^2\\) and \\\\\\[y\\].";
198 expect(formulas(markdown)).toEqual([
199 ["inline", "x^2"],
200 ["display", "y"],
201 ]);
202 expect(text(markdown)).toBe("Path \\ and \\.");
203 });
204
205 it.each([
206 ["$a`$HOME`", "$HOME"],
207 ["\\(a`\\)`", "\\)"],
208 ["\\(a\\alpha`\\)`", "\\)"],
209 ["text \\[a`\\]`", "\\]"],
210 ])("preserves an adjacent code span in %s", (source, code) => {
211 expect(formulas(source)).toEqual([]);
212 expect(collect(parse(source)).find((node) => node.type === "inlineCode")).toMatchObject({
213 value: code,
214 });
215 });
216
217 it.each([1, 2])("respects %i backslashes before a display fence", (count) => {
218 const prefix = "a+" + "\\".repeat(count);
219 for (const opening of ["$$", "$$\n"]) {
220 const source = `${opening}${prefix}$$\nb\n$$`;
221 expect(formulas(source)).toEqual([["display", count % 2 === 1 ? `${prefix}$$\nb` : prefix]]);
222 expect(markdownMathRanges(source)[0]?.source).toBe(
223 count % 2 === 1 ? source : `${opening}${prefix}$$`,
224 );
225 }
226 });
227
228 it("keeps TeX escapes inside formulas", () => {
229 expect(formulas("Cost $a\\$b$ and \\(\\left(x\\right)\\).")).toEqual([
230 ["inline", "a\\$b"],
231 ["inline", "\\left(x\\right)"],
232 ]);
233 });
234
235 it("keeps source offsets of the original message", () => {
236 const markdown = "- [ ] check \\(x\\)\n- [x] then $y$";
237 const offsets = collect(parse(markdown)).flatMap((node) =>
238 node.type === "inlineMath" || node.type === "listItem"
239 ? [[node.type, node.position?.start.offset]]
240 : [],
241 );
242 expect(offsets).toEqual([
243 ["listItem", 0],
244 ["inlineMath", markdown.indexOf("\\(")],
245 ["listItem", markdown.indexOf("- [x]")],
246 ["inlineMath", markdown.indexOf("$y$")],
247 ]);
248 });
249
250 it.each(["\n", "\r\n", "\r"])("keeps authored lines and UTF-16 offsets with %j", (eol) => {
251 const source = `😀 \\(x${eol}+y\\)${eol}${eol}- [ ] check $z$`;
252 const nodes = collect(parse(source));
253 const formula = nodes.find(
254 (node) => node.type === "inlineMath" && node.value === "x" + eol + "+y",
255 );
256 expect(formula?.position?.start).toEqual({ line: 1, column: 4, offset: 3 });
257 expect(formula?.position?.end).toEqual({
258 line: 2,
259 column: 5,
260 offset: source.indexOf(eol + eol),
261 });
262 expect(nodes.find((node) => node.type === "listItem")?.position?.start).toEqual({
263 line: 4,
264 column: 1,
265 offset: source.indexOf("- [ ]"),
266 });
267 });
268
269 it("shows unfinished formulas as text while a response streams", () => {
270 const message = "Energy \\(E = mc^2\\) and $$\\sum_i x_i$$ done.";
271 for (let end = 0; end <= message.length; end++) {
272 expect(() => parse(message.slice(0, end))).not.toThrow();
273 }
274 expect(text("Energy \\(E = mc")).toBe("Energy (E = mc");
275 expect(formulas("$$\nx^2")).toEqual([]);
276 expect(formulas("\\[\nx^2\n- y")).toEqual([]);
277 expect(formulas(message)).toEqual([
278 ["inline", "E = mc^2"],
279 ["display", "\\sum_i x_i"],
280 ]);
281 });
282});
283
284describe("math in GFM tables", () => {
285 it("renders math in link labels while retaining dollars in their destinations", () => {
286 const row = "| [norm $|x|$](https://example.com/$path) | y$ |";
287 for (const source of [`${row}\n| --- | --- |`, `| First | Second |\n| --- | --- |\n${row}`]) {
288 const tree = parse(source);
289 const table = tree.children[0];
290 expect(table?.type).toBe("table");
291 if (table?.type !== "table") throw new Error("Missing table");
292 expect(table.children.every((entry) => entry.children.length === 2)).toBe(true);
293 expect(markdownMathRanges(source, tree).map((range) => range.tex)).toEqual(["|x|"]);
294 expect(collect(tree).find((node) => node.type === "link")).toMatchObject({
295 url: "https://example.com/$path",
296 });
297 }
298 });
299
300 it.each(["[a](https://example.com/$x)", "`$x`", '<span title="$x">a</span>'])(
301 "preserves table separators after literal dollars in %s",
302 (literal) => {
303 const row = `| ${literal} | y$ |`;
304 for (const source of [`${row}\n| --- | --- |`, `| First | Second |\n| --- | --- |\n${row}`]) {
305 const tree = parse(source);
306 const table = tree.children[0];
307 expect(table?.type).toBe("table");
308 if (table?.type !== "table") throw new Error("Missing table");
309 expect(table.children.every((entry) => entry.children.length === 2)).toBe(true);
310 expect(markdownMathRanges(source, tree)).toEqual([]);
311 if (literal.startsWith("[a](https://example.com/$x)")) {
312 expect(collect(tree).find((node) => node.type === "link")).toMatchObject({
313 url: "https://example.com/$x",
314 });
315 }
316 }
317 },
318 );
319
320 it("retains multiline math after a prose pipe outside tables", () => {
321 const source = "a | b: \\(x\n+y\\)";
322 expect(formulas(source)).toEqual([["inline", "x\n+y"]]);
323 expect(markdownMathRanges(source)).toEqual([
324 expect.objectContaining({ source: "\\(x\n+y\\)", inTable: false }),
325 ]);
326 });
327 it.each(["\\(", "\\["])("keeps table rows intact when %s closes in a later row", (opener) => {
328 const closer = opener === "\\(" ? "\\)" : "\\]";
329 const source = `| Name | Value |\n| --- | --- |\n| first | ${opener}x |\n| second | y${closer} |\n| valid | $|z|$ |`;
330 const tree = parseMarkdownWithMath(source);
331 const table = tree.children[0];
332 expect(table?.type).toBe("table");
333 if (table?.type !== "table") throw new Error("Missing table");
334 expect(table.children.map((row) => row.children.length)).toEqual([2, 2, 2, 2]);
335 expect(
336 collect(tree)
337 .filter((node) => node.type === "text")
338 .map((node) => node.value),
339 ).toEqual([
340 "Name",
341 "Value",
342 "first",
343 opener.slice(1) + "x",
344 "second",
345 "y" + closer.slice(1),
346 "valid",
347 ]);
348 expect(markdownMathRanges(source)).toEqual([
349 expect.objectContaining({ source: "$|z|$", tex: "|z|" }),
350 ]);
351 expect(JSON.stringify(tree)).not.toMatch(/[\uE000-\uF8FF]/);
352 });
353 it.each(["\\(\\frac{1}{|x|}\\)", "$$|x|$$", "$|x|$"])(
354 "keeps pipes inside %s in one cell and preserves source offsets",
355 (formula) => {
356 const source = `| Name | Value |\n| --- | --- |\n| x | ${formula} |\n\n- [ ] next`;
357 const tree = parseMarkdownWithMath(source);
358 const table = tree.children[0];
359 expect(table?.type).toBe("table");
360 if (table?.type !== "table") throw new Error("Missing table");
361 expect(table.children.map((row) => row.children.length)).toEqual([2, 2]);
362 expect(markdownMathRanges(source)).toEqual([
363 expect.objectContaining({ source: formula, start: source.indexOf(formula) }),
364 ]);
365 const checkbox = collect(tree).find((node) => node.type === "listItem");
366 expect(checkbox?.position?.start.offset).toBe(source.indexOf("- [ ]"));
367 },
368 );
369
370 it("preserves TeX pipe commands inside and outside tables", () => {
371 const source = "| Formula |\n| --- |\n| \\(\\|x\\|\\) |\n\n\\(\\|x\\|\\)";
372 expect(markdownMathRanges(source).map((formula) => formula.tex)).toEqual([
373 "\\|x\\|",
374 "\\|x\\|",
375 ]);
376 });
377
378 it("leaves code, URLs, incomplete math, and private-use text unchanged", () => {
379 const source =
380 "| Value | Other |\n| --- | --- |\n| `\\(a\\|b\\)` | \\(x|y\\) |\n\n\uE000\n\n[link](https://example.com/$x|y$)\n\nUnclosed \\(x|y";
381 const tree = parseMarkdownWithMath(source);
382 expect(markdownMathRanges(source).map((formula) => formula.tex)).toEqual(["x|y"]);
383 expect(JSON.stringify(tree)).toContain("https://example.com/$x|y$");
384 expect(collect(tree).find((node) => node.type === "inlineCode")).toMatchObject({
385 value: "\\(a|b\\)",
386 });
387 expect(JSON.stringify(tree)).toContain("\uE000");
388 });
389
390 it("keeps a display equation's equality line out of Markdown headings", () => {
391 const source = "\\[\na\n=\nb\n\\]";
392 expect(parseMarkdownWithMath(source).children).toEqual([
393 expect.objectContaining({ type: "math", value: "a\n=\nb" }),
394 ]);
395 });
396});