diff --git a/packages/ui/src/components/chat/markdown/markdown-shiki.worker.ts b/packages/ui/src/components/chat/markdown/markdown-shiki.worker.ts index a06b81fc..bf0891f0 100644 --- a/packages/ui/src/components/chat/markdown/markdown-shiki.worker.ts +++ b/packages/ui/src/components/chat/markdown/markdown-shiki.worker.ts @@ -1,6 +1,10 @@ /// -import { bundledLanguages, createHighlighter, type BundledLanguage, type ThemedToken } from 'shiki'; +import { bundledLanguages, createHighlighter, type BundledLanguage, type LanguageRegistration, type ThemedToken } from 'shiki'; +import { + isTemplateCallLanguageId, + sanitizeTemplateCallGrammar, +} from '../../../lib/shiki/sanitizeTemplateCallGrammar'; import { MARKDOWN_SHIKI_THEME, MARKDOWN_SHIKI_THEME_DEFINITION } from './markdownShikiThemeDefinition'; import type { MarkdownWorkerRequest, MarkdownWorkerResponse } from './markdown-worker-protocol'; @@ -60,11 +64,28 @@ self.onmessage = (event: MessageEvent) => { type Instance = Awaited>; +type BundledLanguageModule = { default: LanguageRegistration[] }; + +/** + * Load a language, neutralizing the catastrophic JS/TS `template-call` rule + * before it reaches the Oniguruma scanner (see sanitizeTemplateCallGrammar). + */ +const loadLanguageSafe = async (instance: Instance, lang: BundledLanguage): Promise => { + if (!isTemplateCallLanguageId(lang)) { + await instance.loadLanguage(bundledLanguages[lang]); + return; + } + + const mod = (await bundledLanguages[lang]()) as BundledLanguageModule; + const grammars = mod.default.map((grammar) => sanitizeTemplateCallGrammar(grammar)); + await instance.loadLanguage(...grammars); +}; + const resolveLanguage = async (instance: Instance, requested: string): Promise => { let lang = requested in bundledLanguages ? requested : 'text'; if (lang !== 'text' && !instance.getLoadedLanguages().includes(lang)) { try { - await instance.loadLanguage(bundledLanguages[lang as BundledLanguage]); + await loadLanguageSafe(instance, lang as BundledLanguage); } catch { lang = 'text'; } diff --git a/packages/ui/src/components/chat/markdown/markdown-worker-timeout.ts b/packages/ui/src/components/chat/markdown/markdown-worker-timeout.ts new file mode 100644 index 00000000..1e82763c --- /dev/null +++ b/packages/ui/src/components/chat/markdown/markdown-worker-timeout.ts @@ -0,0 +1,6 @@ +/** + * Safety-net budget for a single Shiki worker tokenize request. + * Healthy files finish well under this; catastrophic Oniguruma backtracking + * must not run unbounded (openchamber/openchamber#2587). + */ +export const HIGHLIGHT_REQUEST_TIMEOUT_MS = 5_000; diff --git a/packages/ui/src/components/chat/markdown/markdown-worker.hang.test.ts b/packages/ui/src/components/chat/markdown/markdown-worker.hang.test.ts new file mode 100644 index 00000000..df591b2c --- /dev/null +++ b/packages/ui/src/components/chat/markdown/markdown-worker.hang.test.ts @@ -0,0 +1,12 @@ +import { describe, expect, test } from 'bun:test'; + +import { HIGHLIGHT_REQUEST_TIMEOUT_MS } from './markdown-worker-timeout'; + +describe('markdown-worker hang safety', () => { + test('exposes a finite highlight timeout budget', () => { + // Catastrophic Oniguruma backtracking must not run unbounded; the main + // thread terminates the worker after this budget (openchamber/openchamber#2587). + expect(HIGHLIGHT_REQUEST_TIMEOUT_MS).toBeGreaterThan(0); + expect(HIGHLIGHT_REQUEST_TIMEOUT_MS).toBeLessThan(15_001); + }); +}); diff --git a/packages/ui/src/components/chat/markdown/markdown-worker.ts b/packages/ui/src/components/chat/markdown/markdown-worker.ts index 85d061bd..b97372e3 100644 --- a/packages/ui/src/components/chat/markdown/markdown-worker.ts +++ b/packages/ui/src/components/chat/markdown/markdown-worker.ts @@ -1,23 +1,42 @@ import MarkdownShikiWorkerUrl from './markdown-shiki.worker.ts?worker&url'; import type { MarkdownTokenRun, MarkdownWorkerRequest, MarkdownWorkerResponse } from './markdown-worker-protocol'; +import { HIGHLIGHT_REQUEST_TIMEOUT_MS } from './markdown-worker-timeout'; -// Main-thread client for the markdown Shiki worker. Moves syntax tokenization +export { HIGHLIGHT_REQUEST_TIMEOUT_MS } from './markdown-worker-timeout'; + +// Main-thread client for the markdown Shiki Web Worker. Moves syntax tokenization // off the UI thread: a closed code block is shipped to the worker, which returns // ready-to-splice Shiki HTML. On any failure (no worker support, worker crash, -// tokenization error) the promise resolves to `null` and the caller keeps the -// escaped plain-text code — highlighting never falls back onto the main thread. +// tokenization error, or hang timeout) the promise resolves to `null` and the +// caller keeps the escaped plain-text code — highlighting never falls back onto +// the main thread. +// +// The timeout exists because TextMate grammars can enter catastrophic backtracking +// on the Oniguruma WASM engine (openchamber/openchamber#2587). Matching is sync +// inside the worker, so the only way to reclaim memory is to terminate it from +// this thread when a request exceeds the budget. type PendingResolver = (response: MarkdownWorkerResponse | null) => void; +type PendingEntry = { + resolve: PendingResolver; + timer: ReturnType; +}; + let worker: Worker | undefined; let nextId = 0; -const pending = new Map(); +const pending = new Map(); // Theme names whose full definition we've already shipped to the live worker, so // repeat tokenization sends only the name (not the whole theme object) again. const sentThemes = new Set(); +const clearPendingTimers = (): void => { + pending.forEach((entry) => clearTimeout(entry.timer)); +}; + const failAll = (): void => { - pending.forEach((resolve) => resolve(null)); + clearPendingTimers(); + pending.forEach((entry) => entry.resolve(null)); pending.clear(); sentThemes.clear(); worker?.terminate(); @@ -34,10 +53,11 @@ const getWorker = (): Worker | undefined => { return undefined; } worker.onmessage = (event: MessageEvent) => { - const resolve = pending.get(event.data.id); - if (!resolve) return; + const entry = pending.get(event.data.id); + if (!entry) return; + clearTimeout(entry.timer); pending.delete(event.data.id); - resolve(event.data); + entry.resolve(event.data); }; worker.onerror = failAll; worker.onmessageerror = failAll; @@ -50,7 +70,14 @@ const request = (payload: (id: number) => MarkdownWorkerRequest): Promise((resolve) => { - pending.set(id, resolve); + const timer = setTimeout(() => { + if (!pending.has(id)) return; + // Hung tokenize (e.g. catastrophic backtracking): kill the worker so the + // WASM heap is freed instead of growing until the renderer OOMs. + console.warn(`Shiki worker highlight timed out after ${HIGHLIGHT_REQUEST_TIMEOUT_MS}ms; terminating worker`); + failAll(); + }, HIGHLIGHT_REQUEST_TIMEOUT_MS); + pending.set(id, { resolve, timer }); instance.postMessage(payload(id)); }); }; diff --git a/packages/ui/src/lib/shiki/sanitizeTemplateCallGrammar.test.ts b/packages/ui/src/lib/shiki/sanitizeTemplateCallGrammar.test.ts new file mode 100644 index 00000000..bc10ceaf --- /dev/null +++ b/packages/ui/src/lib/shiki/sanitizeTemplateCallGrammar.test.ts @@ -0,0 +1,75 @@ +import { describe, expect, test } from 'bun:test'; +import { bundledLanguages, createHighlighter, type LanguageRegistration } from 'shiki'; + +import { + hasCatastrophicTemplateCall, + isTemplateCallLanguageId, + sanitizeTemplateCallGrammar, + TEMPLATE_CALL_LANGUAGE_IDS, +} from './sanitizeTemplateCallGrammar'; + +type BundledLanguageModule = { default: LanguageRegistration[] }; + +const loadBundledGrammar = async (id: (typeof TEMPLATE_CALL_LANGUAGE_IDS)[number]): Promise => { + const mod = (await bundledLanguages[id]()) as BundledLanguageModule; + return mod.default[0]; +}; + +describe('sanitizeTemplateCallGrammar', () => { + test('detects template-call on bundled JS/TS grammars', async () => { + for (const id of TEMPLATE_CALL_LANGUAGE_IDS) { + const grammar = await loadBundledGrammar(id); + expect(isTemplateCallLanguageId(id)).toBe(true); + expect(hasCatastrophicTemplateCall(grammar)).toBe(true); + } + }); + + test('clears template-call patterns without dropping the repository key', async () => { + const grammar = await loadBundledGrammar('javascript'); + const patched = sanitizeTemplateCallGrammar(grammar); + + expect(hasCatastrophicTemplateCall(patched)).toBe(false); + expect(patched.repository?.['template-call']).toEqual({ patterns: [] }); + // Original left intact (structured clone / spread, not mutate-in-place). + expect(hasCatastrophicTemplateCall(grammar)).toBe(true); + }); + + test('is a no-op when template-call is already empty', () => { + const grammar = { + name: 'javascript', + scopeName: 'source.js', + patterns: [], + repository: { 'template-call': { patterns: [] } }, + } satisfies LanguageRegistration; + expect(sanitizeTemplateCallGrammar(grammar)).toBe(grammar); + }); + + test('highlights template-literal fixtures within a tight budget after sanitize', async () => { + const mod = (await bundledLanguages.javascript()) as BundledLanguageModule; + const patched = mod.default.map((grammar) => sanitizeTemplateCallGrammar(grammar)); + + const highlighter = await createHighlighter({ + themes: ['github-dark'], + langs: patched, + }); + + // Representative content from openchamber/openchamber#2587, scaled to ~14KB. + const fixture = `const snapshot = { source: \`\${session.source}\`, fetchedAt: \`\${Date.now()}\` }; +const label = \`Account \${index + 1}\`; +function render(account) { + return html\`
\${account.name}
\`; +} +`.repeat(80); + + expect(fixture.length).toBeGreaterThan(10_000); + + const started = performance.now(); + const html = highlighter.codeToHtml(fixture, { lang: 'javascript', theme: 'github-dark' }); + const elapsedMs = performance.now() - started; + highlighter.dispose(); + + expect(html.length).toBeGreaterThan(0); + // Catastrophic backtracking hangs for seconds–minutes; healthy tokenize is well under 1s. + expect(elapsedMs).toBeLessThan(2_000); + }); +}); \ No newline at end of file diff --git a/packages/ui/src/lib/shiki/sanitizeTemplateCallGrammar.ts b/packages/ui/src/lib/shiki/sanitizeTemplateCallGrammar.ts new file mode 100644 index 00000000..dafb60de --- /dev/null +++ b/packages/ui/src/lib/shiki/sanitizeTemplateCallGrammar.ts @@ -0,0 +1,45 @@ +/** + * Neutralize the JavaScript/TypeScript TextMate `template-call` rule. + * + * Upstream grammars use a triple-nested `{()[]}` lookahead to detect tagged + * templates with type arguments (`foo\`...\``). On the Oniguruma WASM engine + * shipped with Shiki — which does not expose `setRetryLimit` / match-stack + * limits — that pattern can enter exponential backtracking on ordinary + * backtick template literals, grow the WASM heap without bound, and OOM the + * renderer (openchamber/openchamber#2587). + * + * Clearing `template-call` is safe: the plain `#template` rule still highlights + * backticks and simple tagged templates. Only the rare `ident\`...\`` + * form loses its specialized type-argument coloring and falls through to + * normal tokenization. + */ + +type GrammarRepository = Record; + +export type TemplateCallGrammar = { + name?: string; + repository?: GrammarRepository; +}; + +const TEMPLATE_CALL_KEY = 'template-call'; + +export const hasCatastrophicTemplateCall = (grammar: TemplateCallGrammar): boolean => { + const patterns = grammar.repository?.[TEMPLATE_CALL_KEY]?.patterns; + return Array.isArray(patterns) && patterns.length > 0; +}; + +export const sanitizeTemplateCallGrammar = (grammar: T): T => { + if (!hasCatastrophicTemplateCall(grammar)) return grammar; + + const repository = { ...grammar.repository }; + repository[TEMPLATE_CALL_KEY] = { patterns: [] }; + return { ...grammar, repository }; +}; + +/** Language ids whose bundled grammars ship the catastrophic `template-call` rule. */ +export const TEMPLATE_CALL_LANGUAGE_IDS = ['javascript', 'typescript', 'jsx', 'tsx'] as const; + +export type TemplateCallLanguageId = (typeof TEMPLATE_CALL_LANGUAGE_IDS)[number]; + +export const isTemplateCallLanguageId = (lang: string): lang is TemplateCallLanguageId => + (TEMPLATE_CALL_LANGUAGE_IDS as readonly string[]).includes(lang);