Fix display math parsing and whitespace handling in markdown renderer

src/lib/components/chat/Messages/Markdown/MarkdownInlineTokens.svelte

@@ -109,7 +109,7 @@
 		<del><svelte:self id={`${id}-del`} tokens={token.tokens} {onSourceClick} /></del>
 	{:else if token.type === 'inlineKatex'}
 		{#if token.text}
-			<KatexRenderer content={token.text} displayMode={false} />
+			<KatexRenderer content={token.text} displayMode={token?.displayMode ?? false} />
 		{/if}
 	{:else if token.type === 'iframe'}
 		<iframe

src/lib/utils/marked/katex-extension.test.ts

@@ -0,0 +1,115 @@
+import { describe, expect, it } from 'vitest';
+import { Marked } from 'marked';
+
+import katexExtension from './katex-extension';
+
+const createMarked = () => {
+	const parser = new Marked();
+	parser.use(katexExtension({ throwOnError: false, breaks: true }));
+	return parser;
+};
+
+const collectKatexTokens = (tokens: any[]): any[] => {
+	const result: any[] = [];
+
+	for (const token of tokens) {
+		if (token.type === 'inlineKatex' || token.type === 'blockKatex') {
+			result.push(token);
+		}
+
+		if (token.tokens) {
+			result.push(...collectKatexTokens(token.tokens));
+		}
+	}
+
+	return result;
+};
+
+describe('katex-extension', () => {
+	it('renders balanced single-line $$...$$ as display math', () => {
+		const tokens = createMarked().lexer('$$x^2$$');
+
+		expect(tokens).toHaveLength(1);
+		expect(tokens[0]).toMatchObject({
+			type: 'blockKatex',
+			text: 'x^2',
+			displayMode: true
+		});
+	});
+
+	it('renders multiline $$...$$ as display math', () => {
+		const input = '$$\n\\int_0^1 x^2 \\, dx\n= \\frac{1}{3}\n$$';
+		const tokens = createMarked().lexer(input);
+
+		expect(tokens).toHaveLength(1);
+		expect(tokens[0]).toMatchObject({
+			type: 'blockKatex',
+			text: '\n\\int_0^1 x^2 \\, dx\n= \\frac{1}{3}\n',
+			displayMode: true
+		});
+	});
+
+	it('treats trailing spaces before newline as a block boundary', () => {
+		const tokens = createMarked().lexer('$$x^2$$   \nnext line');
+
+		expect(tokens[0]).toMatchObject({
+			type: 'blockKatex',
+			text: 'x^2',
+			displayMode: true
+		});
+	});
+
+	it('treats CRLF after closing $$ as a block boundary', () => {
+		const tokens = createMarked().lexer('$$x^2$$\r\nnext line');
+
+		expect(tokens[0]).toMatchObject({
+			type: 'blockKatex',
+			text: 'x^2',
+			displayMode: true
+		});
+	});
+
+	it('leaves stray $$ as plain text', () => {
+		const tokens = createMarked().lexer('stray $$ delimiter');
+
+		expect(collectKatexTokens(tokens)).toHaveLength(0);
+		expect(tokens[0]?.type).toBe('paragraph');
+		expect(tokens[0]?.text).toContain('stray');
+		expect(tokens[0]?.text).toContain('$$ delimiter');
+	});
+
+	it('leaves normal prose with dollar signs untouched', () => {
+		const tokens = createMarked().lexer('The total is $5.00 in normal prose.');
+
+		expect(collectKatexTokens(tokens)).toHaveLength(0);
+		expect(tokens[0]).toMatchObject({
+			type: 'paragraph',
+			text: 'The total is $5.00 in normal prose.'
+		});
+	});
+
+	it('keeps punctuation after closing $$ attached to the surrounding paragraph', () => {
+		const tokens = createMarked().lexer('Result: $$x^2$$.');
+		const mathTokens = collectKatexTokens(tokens);
+
+		expect(tokens[0]?.type).toBe('paragraph');
+		expect(mathTokens).toHaveLength(1);
+		expect(mathTokens[0]).toMatchObject({
+			type: 'inlineKatex',
+			text: 'x^2',
+			displayMode: true
+		});
+		expect(tokens[0]?.text).toContain('.');
+	});
+
+	it('does not interfere with fenced math code blocks', () => {
+		const tokens = createMarked().lexer('```math\nx^2\n```');
+
+		expect(tokens).toHaveLength(1);
+		expect(tokens[0]).toMatchObject({
+			type: 'code',
+			lang: 'math',
+			text: 'x^2'
+		});
+	});
+});

src/lib/utils/marked/katex-extension.ts

@@ -66,6 +66,62 @@ function generateRegexRules(delimiters) {
 
 const { inlineRule, blockRule } = generateRegexRules(DELIMITER_LIST);
 
+function hasAllowedTrailingChar(src: string, index: number) {
+	return index >= src.length || ALLOWED_SURROUNDING_CHARS_REGEX.test(src.charAt(index));
+}
+
+function hasBlockTrailingChar(src: string, index: number) {
+	return /^(?:[ \t]*\r?\n|$)/.test(src.slice(index));
+}
+
+function findDisplayMathEnd(src: string, startIndex: number) {
+	for (let i = startIndex; i < src.length - 1; i++) {
+		if (src.charAt(i) === '\\') {
+			i += 1;
+			continue;
+		}
+
+		if (src.charAt(i) === '$' && src.charAt(i + 1) === '$') {
+			return i;
+		}
+	}
+
+	return -1;
+}
+
+export function tokenizeDisplayMath(
+	src: string,
+	type: 'inlineKatex' | 'blockKatex',
+	requireBlockBoundary = false
+) {
+	if (!src.startsWith('$$')) {
+		return;
+	}
+
+	const endIndex = findDisplayMathEnd(src, 2);
+	if (endIndex === -1) {
+		return;
+	}
+
+	const raw = src.slice(0, endIndex + 2);
+	const text = raw.slice(2, -2);
+
+	if (!text.trim() || !hasAllowedTrailingChar(src, endIndex + 2)) {
+		return;
+	}
+
+	if (requireBlockBoundary && !hasBlockTrailingChar(src, endIndex + 2)) {
+		return;
+	}
+
+	return {
+		type,
+		raw,
+		text,
+		displayMode: true
+	};
+}
+
 export default function (options = {}) {
 	return {
 		extensions: [inlineKatex(options), blockKatex(options)]
@@ -102,6 +158,17 @@ function katexStart(src, displayMode: boolean) {
 }
 
 function katexTokenizer(src, tokens, displayMode: boolean) {
+	if (src.startsWith('$$')) {
+		const displayToken = tokenizeDisplayMath(
+			src,
+			displayMode ? 'blockKatex' : 'inlineKatex',
+			displayMode
+		);
+		if (displayToken) {
+			return displayToken;
+		}
+	}
+
 	const ruleReg = displayMode ? blockRule : inlineRule;
 	const type = displayMode ? 'blockKatex' : 'inlineKatex';