From 15b89b9218b7d2c7239c579aa3d23c2892227ac6 Mon Sep 17 00:00:00 2001
From: Timothy Jaeryang Baek <tim@openwebui.com>
Date: Sun, 12 Apr 2026 16:56:00 -0500
Subject: [PATCH] refac

---
 .../Markdown/MarkdownInlineTokens.svelte      |  2 +-
 src/lib/utils/marked/katex-extension.ts       | 51 +++++++++++++++++++
 2 files changed, 52 insertions(+), 1 deletion(-)

diff --git a/src/lib/components/chat/Messages/Markdown/MarkdownInlineTokens.svelte b/src/lib/components/chat/Messages/Markdown/MarkdownInlineTokens.svelte
index c2c7d81e6..0daf389ee 100644
--- a/src/lib/components/chat/Messages/Markdown/MarkdownInlineTokens.svelte
+++ b/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
diff --git a/src/lib/utils/marked/katex-extension.ts b/src/lib/utils/marked/katex-extension.ts
index 13860fab3..536635f1a 100644
--- a/src/lib/utils/marked/katex-extension.ts
+++ b/src/lib/utils/marked/katex-extension.ts
@@ -66,6 +66,46 @@ function generateRegexRules(delimiters) {
 
 const { inlineRule, blockRule } = generateRegexRules(DELIMITER_LIST);
 
+const isAllowedTrailing = (src: string, i: number): boolean =>
+	i >= src.length || ALLOWED_SURROUNDING_CHARS_REGEX.test(src.charAt(i));
+
+const isBlockBoundary = (src: string, i: number): boolean =>
+	/^(?:[ \t]*\r?\n|$)/.test(src.slice(i));
+
+const findClosingDelimiter = (src: string, i: number): number =>
+	i >= src.length - 1
+		? -1
+		: src[i] === '\\'
+			? findClosingDelimiter(src, i + 2)
+			: src[i] === '$' && src[i + 1] === '$'
+				? i
+				: findClosingDelimiter(src, i + 1);
+
+export const tokenizeDisplayMath = (
+	src: string,
+	type: 'inlineKatex' | 'blockKatex',
+	requireBlockBoundary = false
+) => {
+	if (!src.startsWith('$$')) return;
+
+	const endIndex = findClosingDelimiter(src, 2);
+	if (endIndex === -1) return;
+
+	const raw = src.slice(0, endIndex + 2);
+	const text = raw.slice(2, -2);
+	const afterClose = endIndex + 2;
+
+	const validators: Array<() => boolean> = [
+		() => text.trim().length > 0,
+		() => isAllowedTrailing(src, afterClose),
+		() => !requireBlockBoundary || isBlockBoundary(src, afterClose)
+	];
+
+	return validators.every((v) => v())
+		? { type, raw, text, displayMode: true }
+		: undefined;
+};
+
 export default function (options = {}) {
 	return {
 		extensions: [inlineKatex(options), blockKatex(options)]
@@ -102,6 +142,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';