diff --git a/packages/shared/src/markdownMath.test.ts b/packages/shared/src/markdownMath.test.ts index e34dde0e9e212..d37f36e9445ec 100644 --- a/packages/shared/src/markdownMath.test.ts +++ b/packages/shared/src/markdownMath.test.ts @@ -1,4 +1,8 @@ import type { Nodes } from "mdast"; +import remarkGfm from "remark-gfm"; +import remarkMath from "remark-math"; +import remarkParse from "remark-parse"; +import { unified } from "unified"; import { describe, expect, it } from "vite-plus/test"; import { mathKind } from "./markdownMath.ts"; @@ -266,6 +270,63 @@ describe("chat math parsing", () => { }); }); + it.each(["\\( a ", "\\[ a ", "\\[\n"])( + "bounds tokenizer work for repeated unclosed %j", + (opener) => { + let consumed = 0; + const parser = unified() + .use(remarkParse) + .use(remarkGfm) + .use(remarkMath) + .use(function () { + for (const extension of this.data().micromarkExtensions ?? []) { + for (const constructs of [extension.text, extension.flow]) { + const entries = constructs?.[92]; + if (!entries || !constructs) continue; + constructs[92] = (Array.isArray(entries) ? entries : [entries]).map((construct) => ({ + ...construct, + tokenize(effects, ok, nok) { + return construct.tokenize.call( + this, + { + ...effects, + consume(code) { + consumed++; + effects.consume(code); + }, + }, + ok, + nok, + ); + }, + })); + } + } + }); + const source = opener.repeat(200); + expect( + collect(parser.parse(source)).some((node) => /^(inlineMath|math)$/.test(node.type)), + ).toBe(false); + expect(consumed).toBeGreaterThan(source.length); + expect(consumed).toBeLessThan(source.length * 24); + }, + ); + + it.each([ + ["\\( a \\[b\\]", [["display", "b"]]], + ["\\[ a \\(b\\)", [["inline", "b"]]], + ["\\( a \\( b`code` \\(c\\)", [["inline", "c"]]], + ["\\[1\\] \\[x\\]", [["display", "x"]]], + ["\\( a\n\n\\(b\\)", [["inline", "b"]]], + ["> \\[\n> \\[\n\n\\[\nx\n\\]", [["display", "x"]]], + ["\\[\n\\[\nx\n\\]", [["display", "\\[\nx"]]], + ["\\(a \\\\\\(b\\)", [["inline", "a \\\\\\(b"]]], + ["\\(a \\\\) \\(b\\)", [["inline", "a \\\\) \\(b"]]], + ["| a | b |\n| --- | --- |\n| \\(a \\(b | c |\n\n\\(x\\)", [["inline", "x"]]], + ])("keeps valid formulas after failed scans in %j", (source, expected) => { + expect(formulas(source)).toEqual(expected); + }); + it("shows unfinished formulas as text while a response streams", () => { const message = "Energy \\(E = mc^2\\) and $$\\sum_i x_i$$ done."; for (let end = 0; end <= message.length; end++) { diff --git a/patches/micromark-extension-llm-math@3.1.0.patch b/patches/micromark-extension-llm-math@3.1.0.patch index e09a8d4d889eb..c12fad9b9bcf8 100644 --- a/patches/micromark-extension-llm-math@3.1.0.patch +++ b/patches/micromark-extension-llm-math@3.1.0.patch @@ -163,7 +163,18 @@ diff --git a/dev/lib/math-flow.js b/dev/lib/math-flow.js diff --git a/dev/lib/math-tex-flow.js b/dev/lib/math-tex-flow.js --- a/dev/lib/math-tex-flow.js +++ b/dev/lib/math-tex-flow.js -@@ -26,6 +26,7 @@ +@@ -7,6 +7,10 @@ + import {markdownLineEnding} from 'micromark-util-character' + import {codes, constants, types} from 'micromark-util-symbol' + ++// Only EOF proves that every later opener in this flow lacks a closing fence. ++/** @type {WeakMap} */ ++const failedUntil = new WeakMap() ++ + /** @type {Construct} */ + export const mathTexFlow = { + tokenize: tokenizeMathFenced, +@@ -26,6 +30,7 @@ */ function tokenizeMathFenced(effects, ok, nok) { const self = this @@ -171,7 +182,15 @@ diff --git a/dev/lib/math-tex-flow.js b/dev/lib/math-tex-flow.js const tail = self.events[self.events.length - 1] const initialSize = tail && tail[1].type === types.linePrefix -@@ -90,6 +91,7 @@ +@@ -47,6 +52,7 @@ + * @type {State} + */ + function start(code) { ++ if ((failedUntil.get(self) || 0) > self.now().offset) return nok(code) + assert(code === codes.backslash, 'expected `\\`') + effects.enter('mathFlow') + effects.enter('mathFlowFence') +@@ -90,6 +96,7 @@ * @type {State} */ function metaAfter(code) { @@ -179,15 +198,18 @@ diff --git a/dev/lib/math-tex-flow.js b/dev/lib/math-tex-flow.js effects.exit('mathFlowFence') if (self.interrupt) { -@@ -245,6 +247,7 @@ +@@ -245,6 +252,10 @@ * @type {State} */ function after(code) { -+ if (!closed) return nok(code) ++ if (!closed) { ++ if (code === codes.eof) failedUntil.set(self, self.now().offset) ++ return nok(code) ++ } effects.exit('mathFlow') return ok(code) } -@@ -327,6 +330,7 @@ +@@ -327,6 +338,7 @@ function afterSequenceClose(code) { if (code === codes.eof || markdownLineEnding(code)) { effects.exit('mathFlowFence') @@ -198,7 +220,32 @@ diff --git a/dev/lib/math-tex-flow.js b/dev/lib/math-tex-flow.js diff --git a/dev/lib/math-tex-text.js b/dev/lib/math-tex-text.js --- a/dev/lib/math-tex-text.js +++ b/dev/lib/math-tex-text.js -@@ -25,6 +25,7 @@ +@@ -1,5 +1,5 @@ + /** +- * @import {Construct, Previous, State, Token, TokenizeContext, Tokenizer} from 'micromark-util-types' ++ * @import {Construct, Event, Previous, State, Token, TokenizeContext, Tokenizer} from 'micromark-util-types' + */ + + import {ok as assert} from 'devlop' +@@ -7,6 +7,12 @@ + import {codes, types} from 'micromark-util-symbol' + import {mathText} from './math-text.js' + ++// Failed scans are local to a text tokenizer and delimiter, never shared across documents. ++/** @type {WeakMap, Record>} */ ++const failedUntil = new WeakMap() ++/** @type {WeakMap, Record>} */ ++const singleLineFailedUntil = new WeakMap() ++ + /** @type {Construct} */ + export const mathTexText = { + tokenize: tokenizeMathText, +@@ -21,10 +27,12 @@ + */ + function tokenizeMathText(effects, ok, nok) { + const self = this ++ const failures = self.mathSingleLine ? singleLineFailedUntil : failedUntil + /** @type {Token} */ let token /** @type {number} */ let openSequenceCode @@ -206,25 +253,50 @@ diff --git a/dev/lib/math-tex-text.js b/dev/lib/math-tex-text.js return start -@@ -83,7 +84,7 @@ +@@ -66,6 +74,8 @@ + return nok(code) + } + ++ if ((failures.get(self.events)?.[code] || 0) > self.now().offset) return nok(code) ++ + openSequenceCode = code + effects.consume(code) + effects.exit('mathTextSequence') +@@ -83,8 +93,8 @@ * @type {State} */ function between(code) { - if (code === codes.eof) { +- return nok(code) + if (code === codes.eof || code === codes.graveAccent || (self.mathSingleLine && markdownLineEnding(code))) { - return nok(code) ++ return unmatched(code) } -@@ -124,6 +125,8 @@ + if (code === codes.backslash) { +@@ -124,6 +134,8 @@ * @type {State} */ function data(code) { -+ if (code === codes.graveAccent) return nok(code) ++ if (code === codes.graveAccent) return unmatched(code) + if (code !== codes.backslash && code !== codes.eof && code !== codes.space && code !== codes.comma && !markdownLineEnding(code) && !(code >= codes.digit0 && code <= codes.digit9)) substantive = true if ( code === codes.eof || code === codes.space || -@@ -155,6 +158,7 @@ +@@ -138,6 +150,13 @@ + return data + } + ++ function unmatched(code) { ++ const ends = failures.get(self.events) || {} ++ ends[openSequenceCode] = self.now().offset ++ failures.set(self.events, ends) ++ return nok(code) ++ } ++ + /** + * In closing sequence. + * +@@ -155,6 +174,7 @@ ? codes.rightParenthesis : codes.rightSquareBracket if (code === closeSequenceCode) { @@ -232,7 +304,7 @@ diff --git a/dev/lib/math-tex-text.js b/dev/lib/math-tex-text.js effects.consume(code) effects.exit('mathTextSequence') effects.exit('mathText') -@@ -180,5 +184,8 @@ +@@ -180,5 +200,8 @@ */ function previous(code) { // If there is a previous code, there will always be a tail. diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 6451f3f97cb29..5c7a690bd85de 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -377,7 +377,7 @@ patchedDependencies: mdast-util-math@3.0.0: afe6b2c89e1eaddfd3c90cf2723f12c9a17d513e65beff51653781486013251c mdast-util-to-hast@13.2.1: 57ea46c73aea3c876ab8d351fe0482ddd081dcb4a0610576e399dbe5eb59d2b1 micromark-extension-gfm-table@2.1.1: 9ad23a530e45d9f83cba1da8654fe16d5af8b59d5c73d5c7a7c91fd25b9be1d8 - micromark-extension-llm-math@3.1.0: 4eeb9d9e27ebc75e98795e21c0719005483dad7d233055698257f7d10e91603b + micromark-extension-llm-math@3.1.0: 6c191d17ea797a4ac6cc88b901041c5e6c5d4c4e508102a033e3f14e95bd2add micromark@4.0.2: cb79edde7eca83bbcaaafb6ceb7b044b81b987314cde4910823b15257ecada67 node-pty@1.2.0-beta.15: f1e680afec85f650699ea756f91f469844ec4c976f4f168492a45c6b3f67d651 react-native-gesture-handler@3.5.0: 171379dd378f8f11e9187876dec2116e5cbf5f21cef6576cdde6956518324f98 @@ -20458,7 +20458,7 @@ snapshots: micromark-util-combine-extensions: 2.0.1 micromark-util-types: 2.0.2 - micromark-extension-llm-math@3.1.0(patch_hash=4eeb9d9e27ebc75e98795e21c0719005483dad7d233055698257f7d10e91603b): + micromark-extension-llm-math@3.1.0(patch_hash=6c191d17ea797a4ac6cc88b901041c5e6c5d4c4e508102a033e3f14e95bd2add): dependencies: '@types/katex': 0.16.8 devlop: 1.1.0 @@ -21787,7 +21787,7 @@ snapshots: dependencies: '@types/mdast': 4.0.4 mdast-util-math: 3.0.0(patch_hash=afe6b2c89e1eaddfd3c90cf2723f12c9a17d513e65beff51653781486013251c)(supports-color@8.1.1) - micromark-extension-math: micromark-extension-llm-math@3.1.0(patch_hash=4eeb9d9e27ebc75e98795e21c0719005483dad7d233055698257f7d10e91603b) + micromark-extension-math: micromark-extension-llm-math@3.1.0(patch_hash=6c191d17ea797a4ac6cc88b901041c5e6c5d4c4e508102a033e3f14e95bd2add) unified: 11.0.5 transitivePeerDependencies: - supports-color