Lazy-load KaTeX and highlight.js, cut per-token markdown render costs

Author: gary149

src/lib/components/CodeBlock.svelte

@@ -15,6 +15,20 @@
 
 	let previewOpen = $state(false);
 
+	// `code` always comes from our own highlighter (hljs or escapeHTML in
+	// marked.ts), which only emits escaped text inside hljs <span> wrappers.
+	// While the fence is still streaming (`loading`), the block re-renders on
+	// every flush, and re-sanitizing the whole growing block each time is the
+	// main-thread hot path of code streaming. Skip DOMPurify during that
+	// window only while the html verifiably matches the highlighter's output
+	// alphabet (raw `<` may open nothing but a span tag): any other markup —
+	// which our highlighter cannot produce — falls back to a full sanitize.
+	// Every completed block still gets sanitized as defense in depth.
+	const NON_HIGHLIGHTER_TAG = /<(?!\/?span[\s>])/i;
+	let sanitizedCode = $derived(
+		loading && !NON_HIGHLIGHTER_TAG.test(code) ? code : DOMPurify.sanitize(code)
+	);
+
 	function hasStrictHtml5Doctype(input: string): boolean {
 		if (!input) return false;
 		const withoutBOM = input.replace(/^\uFEFF/, "");
@@ -64,7 +78,7 @@
 		</div>
 	</div>
 	<pre class="scrollbar-custom overflow-auto px-5 font-mono transition-[height]"><code
-			><!-- eslint-disable svelte/no-at-html-tags -->{@html DOMPurify.sanitize(code)}</code
+			><!-- eslint-disable svelte/no-at-html-tags -->{@html sanitizedCode}</code
 		></pre>
 
 	{#if previewOpen}

src/lib/components/chat/ArtifactPanel.svelte

@@ -7,7 +7,7 @@
 	import { diffLines, diffStats, renderDiffHtml } from "$lib/utils/artifactDiff";
 	import { buildArtifactSrcdoc, isDeployableKind } from "$lib/utils/previewSrcdoc";
 	import { parseExternalUrl } from "$lib/utils/externalLink";
-	import { highlightCode } from "$lib/utils/marked";
+	import { escapeHTML } from "$lib/utils/markedLight";
 	import { artifactPanel, ARTIFACT_PANEL_DEFAULT_FRACTION } from "$lib/stores/artifactPanel.svelte";
 	import { pendingChatInput } from "$lib/stores/pendingChatInput";
 	import { usePublicConfig } from "$lib/utils/PublicConfig.svelte";
@@ -111,8 +111,27 @@
 		}
 	}
 
+	// The highlighter lives in the KaTeX/highlight.js chunk, which is kept out
+	// of the entry bundle (see markedLight.ts). Load it the first time the
+	// panel opens; until it resolves, code renders as escaped plain text and
+	// the effect below re-runs automatically once the real highlighter lands.
+	let highlightCode: ((text: string, lang?: string) => string) | undefined = $state();
+	$effect(() => {
+		if (!artifactPanel.open || highlightCode) return;
+		import("$lib/utils/marked")
+			.then((markedModule) => {
+				highlightCode = markedModule.highlightCode;
+			})
+			.catch(() => {
+				// Chunk failed to load (offline, deploy rotation): keep escaped text
+				// permanently rather than retrying on every effect re-run.
+				highlightCode = (text: string) => escapeHTML(text);
+			});
+	});
+
 	$effect(() => {
 		if (!artifactPanel.open) return;
+		const highlight = highlightCode ?? ((text: string) => escapeHTML(text));
 		// A pending throttled run would paint stale content over whatever the
 		// branches below decide to show
 		clearTimeout(highlightTimer);
@@ -127,9 +146,7 @@
 			// The highlighter runs on the full old/new contents so token colors
 			// survive in the diff (multi-line constructs included).
 			const lang = hljsLanguageFor(version);
-			highlightedCode = DOMPurify.sanitize(
-				renderDiffHtml(diff, (text) => highlightCode(text, lang))
-			);
+			highlightedCode = DOMPurify.sanitize(renderDiffHtml(diff, (text) => highlight(text, lang)));
 			lastHighlightAt = Date.now();
 			return;
 		}
@@ -138,7 +155,7 @@
 		const complete = version.complete;
 
 		const run = () => {
-			highlightedCode = DOMPurify.sanitize(highlightCode(content, lang));
+			highlightedCode = DOMPurify.sanitize(highlight(content, lang));
 			lastHighlightAt = Date.now();
 		};
 

src/lib/components/chat/MarkdownBlock.svelte

@@ -1,5 +1,5 @@
 <script lang="ts">
-	import type { Token } from "$lib/utils/marked";
+	import type { Token } from "$lib/utils/markedLight";
 	import CodeBlock from "../CodeBlock.svelte";
 
 	interface Props {

src/lib/components/chat/MarkdownRenderer.svelte

@@ -1,5 +1,5 @@
 <script lang="ts">
-	import { fallbackBlocks, type BlockToken } from "$lib/utils/marked";
+	import { fallbackBlocks, type BlockToken } from "$lib/utils/markedLight";
 	import {
 		acquireMarkdownClientId,
 		cancelMarkdownClient,

src/lib/utils/markdownWorkerPool.ts

@@ -1,6 +1,20 @@
 import { browser } from "$app/environment";
 import MarkdownWorker from "$lib/workers/markdownWorker?worker";
-import { fallbackBlocks, processBlocks, type BlockToken } from "$lib/utils/marked";
+import { fallbackBlocks, type BlockToken } from "$lib/utils/markedLight";
+
+// The rich pipeline (KaTeX + highlight.js, ~700KB decoded) normally lives only
+// inside the worker's own chunk. The main thread needs it solely when workers
+// are unavailable (no Worker support, strict CSP, repeated worker deaths), so
+// it is loaded on demand and memoized — never in the entry bundle.
+let richModulePromise: Promise<typeof import("$lib/utils/marked")> | undefined;
+const loadRichModule = () =>
+	(richModulePromise ??= import("$lib/utils/marked").catch((err) => {
+		// Never memoize a rejected import: a transient chunk-load failure (deploy
+		// rotation, flaky network) would otherwise disable rich markdown for the
+		// rest of the session. Clearing lets the next render retry.
+		richModulePromise = undefined;
+		throw err;
+	}));
 
 type Source = { title?: string; link: string };
 type ResultCallback = (blocks: BlockToken[], requestId: number) => void;
@@ -72,11 +86,13 @@ function runOnMainThread(job: {
 	requestId: number;
 	onResult: ResultCallback;
 }) {
-	void processBlocks(job.content, job.sources, job.streaming)
+	void loadRichModule()
+		.then((markedModule) => markedModule.processBlocks(job.content, job.sources, job.streaming))
 		.then((blocks) => job.onResult(blocks, job.requestId))
-		// Symmetry with the worker path ("never go silent"): on a parse error still
-		// resolve the render with the lightweight fallback instead of leaking an
-		// unhandled rejection and stranding the renderer on escaped plaintext.
+		// Symmetry with the worker path ("never go silent"): on a parse error (or a
+		// failed chunk load) still resolve the render with the lightweight fallback
+		// instead of leaking an unhandled rejection and stranding the renderer on
+		// escaped plaintext.
 		.catch(() => job.onResult(fallbackBlocks(job.content), job.requestId));
 }
 

src/lib/utils/marked.ts

@@ -31,6 +31,14 @@ import sql from "highlight.js/lib/languages/sql";
 import plaintext from "highlight.js/lib/languages/plaintext";
 import { parseIncompleteMarkdown } from "./parseIncompleteMarkdown";
 import { parseMarkdownIntoBlocks } from "./parseBlocks";
+import { escapeHTML, type BlockToken, type Token } from "./markedLight";
+
+// Re-export the light pieces so existing consumers of this module (the worker,
+// tests) keep a single import site. Main-thread code must import them from
+// ./markedLight directly — importing this module eagerly pulls KaTeX +
+// highlight.js into the entry bundle.
+export { escapeHTML, fallbackBlocks } from "./markedLight";
+export type { BlockToken, CodeToken, TextToken, Token } from "./markedLight";
 
 const bundledLanguages: [string, LanguageFn][] = [
 	["javascript", javascript],
@@ -227,20 +235,6 @@ const katexInlineExtension: TokenizerExtension & RendererExtension = {
 	},
 };
 
-function escapeHTML(content: string) {
-	return content.replace(
-		/[<>&"']/g,
-		(x) =>
-			({
-				"<": "&lt;",
-				">": "&gt;",
-				"&": "&amp;",
-				"'": "&#39;",
-				'"': "&quot;",
-			})[x] || x
-	);
-}
-
 function addInlineCitations(md: string, webSearchSources: SimpleSource[] = []): string {
 	const linkStyle =
 		"color: rgb(59, 130, 246); text-decoration: none; hover:text-decoration: underline;";
@@ -421,19 +415,6 @@ function isFencedBlockClosed(raw?: string): boolean {
 	return closingFencePattern.test(trimmed);
 }
 
-type CodeToken = {
-	type: "code";
-	lang: string;
-	code: string;
-	rawCode: string;
-	isClosed: boolean;
-};
-
-type TextToken = {
-	type: "text";
-	html: string | Promise<string>;
-};
-
 const blockCache = new Map<string, BlockToken>();
 
 // The markdown worker pool keeps a small number of long-lived workers alive for the whole
@@ -499,14 +480,6 @@ export function processTokensSync(content: string, sources: SimpleSource[]): Tok
 	});
 }
 
-export type Token = CodeToken | TextToken;
-
-export type BlockToken = {
-	id: string;
-	content: string;
-	tokens: Token[];
-};
-
 /**
  * Simple hash function for generating stable block IDs
  */
@@ -555,37 +528,6 @@ export async function processBlocks(
 	);
 }
 
-/**
- * Cheap, allocation-light blocks used for SSR and the initial client render.
- *
- * Rich markdown rendering (marked + highlight.js + KaTeX + DOMPurify/jsdom) is
- * intentionally NOT run here: it executes synchronously on the single Node event loop
- * during SSR and, summed across every message of a conversation, can block the loop
- * long enough to fail liveness/readiness health checks (observed event-loop stalls up
- * to ~1.5s). The browser upgrades each message to fully rendered markdown via the
- * markdown worker (or async processBlocks fallback) on mount.
- *
- * The output is deterministic and identical on server and client, so hydration is not
- * affected; only the first paint shows lightly-formatted text before the worker result
- * arrives.
- */
-export function fallbackBlocks(content: string): BlockToken[] {
-	// Static id: it is only used as the {#each} key for this single throwaway block and
-	// has no semantic meaning, so there is no need to hash the (potentially large) content.
-	return [
-		{
-			id: "fallback",
-			content,
-			tokens: [
-				{
-					type: "text",
-					html: `<div style="white-space:pre-wrap;overflow-wrap:anywhere">${escapeHTML(content)}</div>`,
-				},
-			],
-		},
-	];
-}
-
 /**
  * Synchronous version of processBlocks. Kept for non-SSR callers and tests; it is no
  * longer used on the SSR render path (see fallbackBlocks).

src/lib/utils/markedLight.ts

@@ -0,0 +1,72 @@
+// Dependency-free subset of the markdown pipeline.
+//
+// Everything importable from the main-thread entry chunks (MarkdownRenderer,
+// the worker pool, ArtifactPanel) must come from here: the rich pipeline in
+// ./marked.ts pulls in KaTeX + highlight.js (~700KB decoded), which must only
+// ever load inside the markdown worker or through a dynamic import().
+
+export type CodeToken = {
+	type: "code";
+	lang: string;
+	code: string;
+	rawCode: string;
+	isClosed: boolean;
+};
+
+export type TextToken = {
+	type: "text";
+	html: string | Promise<string>;
+};
+
+export type Token = CodeToken | TextToken;
+
+export type BlockToken = {
+	id: string;
+	content: string;
+	tokens: Token[];
+};
+
+export function escapeHTML(content: string) {
+	return content.replace(
+		/[<>&"']/g,
+		(x) =>
+			({
+				"<": "&lt;",
+				">": "&gt;",
+				"&": "&amp;",
+				"'": "&#39;",
+				'"': "&quot;",
+			})[x] || x
+	);
+}
+
+/**
+ * Cheap, allocation-light blocks used for SSR and the initial client render.
+ *
+ * Rich markdown rendering (marked + highlight.js + KaTeX + DOMPurify/jsdom) is
+ * intentionally NOT run here: it executes synchronously on the single Node event loop
+ * during SSR and, summed across every message of a conversation, can block the loop
+ * long enough to fail liveness/readiness health checks (observed event-loop stalls up
+ * to ~1.5s). The browser upgrades each message to fully rendered markdown via the
+ * markdown worker (or async processBlocks fallback) on mount.
+ *
+ * The output is deterministic and identical on server and client, so hydration is not
+ * affected; only the first paint shows lightly-formatted text before the worker result
+ * arrives.
+ */
+export function fallbackBlocks(content: string): BlockToken[] {
+	// Static id: it is only used as the {#each} key for this single throwaway block and
+	// has no semantic meaning, so there is no need to hash the (potentially large) content.
+	return [
+		{
+			id: "fallback",
+			content,
+			tokens: [
+				{
+					type: "text",
+					html: `<div style="white-space:pre-wrap;overflow-wrap:anywhere">${escapeHTML(content)}</div>`,
+				},
+			],
+		},
+	];
+}

src/lib/workers/markdownWorker.ts

@@ -41,17 +41,24 @@ onmessage = async (event) => {
 	}
 
 	try {
-		postMessage(
-			JSON.parse(JSON.stringify({ type: "processed", blocks, requestId })) as OutgoingMessage
-		);
+		// Blocks are plain data (strings/booleans/arrays) and structured-clone
+		// directly; the JSON round-trip below is only needed for exotic values
+		// (e.g. an unresolved promise-valued html) that would fail to clone.
+		postMessage({ type: "processed", blocks, requestId } as OutgoingMessage);
 	} catch {
-		// A block somehow isn't JSON/structured-clone safe. Reply with the fallback so the
-		// pool always gets a reply and the shared worker is never left wedged. fallbackBlocks
-		// output is plain data and always serializes.
-		postMessage({
-			type: "processed",
-			blocks: fallbackBlocks(content),
-			requestId,
-		} as OutgoingMessage);
+		try {
+			postMessage(
+				JSON.parse(JSON.stringify({ type: "processed", blocks, requestId })) as OutgoingMessage
+			);
+		} catch {
+			// A block somehow isn't JSON-safe either. Reply with the fallback so the
+			// pool always gets a reply and the shared worker is never left wedged.
+			// fallbackBlocks output is plain data and always serializes.
+			postMessage({
+				type: "processed",
+				blocks: fallbackBlocks(content),
+				requestId,
+			} as OutgoingMessage);
+		}
 	}
 };