From 464c4ac0caa55bbafd49ff7bf3f723287e17da38 Mon Sep 17 00:00:00 2001 From: Bohdan Triapitsyn Date: Mon, 15 Jun 2026 23:06:54 +0300 Subject: [PATCH] perf(markdown): move code highlighting off the main thread into a Shiki worker MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Tokenize closed code blocks in a dedicated Shiki Web Worker instead of calling the shared highlighter synchronously on the UI thread. This removes the one-shot main-thread highlight stall when a code fence closes on a large block. Streaming behavior is unchanged: the open (streaming) fence still renders as plain text and is highlighted once on close. On any worker failure the block keeps its escaped plain code — highlighting never falls back onto the main thread. - Add markdownShikiThemeDefinition (dependency-free CSS-variable theme) so the worker can use the theme without pulling in @pierre/diffs / React. - Add markdown-worker-protocol, markdown-shiki.worker, and the main-thread markdown-worker client. - Route highlightCodeBlocks through the worker; keep the size/VSCode line guard and mermaid skip on the main thread. - Add shiki as a direct dependency (was transitive via @pierre/diffs). --- bun.lock | 9 +- packages/ui/package.json | 1 + .../chat/markdown/markdown-shiki.worker.ts | 55 +++++++++ .../chat/markdown/markdown-worker-protocol.ts | 13 +++ .../chat/markdown/markdown-worker.ts | 53 +++++++++ .../components/chat/markdown/markdownCore.ts | 43 ++----- .../markdown/markdownShikiThemeDefinition.ts | 100 +++++++++++++++++ .../components/chat/markdown/markdownTheme.ts | 106 +++--------------- 8 files changed, 249 insertions(+), 131 deletions(-) create mode 100644 packages/ui/src/components/chat/markdown/markdown-shiki.worker.ts create mode 100644 packages/ui/src/components/chat/markdown/markdown-worker-protocol.ts create mode 100644 packages/ui/src/components/chat/markdown/markdown-worker.ts create mode 100644 packages/ui/src/components/chat/markdown/markdownShikiThemeDefinition.ts diff --git a/bun.lock b/bun.lock index 580c1ba0..bb9d661f 100644 --- a/bun.lock +++ b/bun.lock @@ -99,7 +99,7 @@ }, "packages/electron": { "name": "@openchamber/electron", - "version": "1.12.4", + "version": "1.13.0", "dependencies": { "@openchamber/web": "workspace:*", "electron-context-menu": "^4.1.2", @@ -114,7 +114,7 @@ }, "packages/ui": { "name": "@openchamber/ui", - "version": "1.12.4", + "version": "1.13.0", "dependencies": { "@base-ui/react": "^1.4.0", "@codemirror/autocomplete": "^6.20.0", @@ -178,6 +178,7 @@ "rehype-katex": "^7.0.1", "remark-math": "^6.0.0", "remend": "^1.2.1", + "shiki": "^3.23.0", "simple-git": "^3.28.0", "sonner": "^2.0.7", "strip-json-comments": "^5.0.3", @@ -216,7 +217,7 @@ }, "packages/vscode": { "name": "openchamber", - "version": "1.12.4", + "version": "1.13.0", "dependencies": { "@openchamber/ui": "workspace:*", "@opencode-ai/sdk": "^1.17.7", @@ -239,7 +240,7 @@ }, "packages/web": { "name": "@openchamber/web", - "version": "1.12.4", + "version": "1.13.0", "bin": { "openchamber": "./bin/cli.js", }, diff --git a/packages/ui/package.json b/packages/ui/package.json index ed9928e0..27789c2e 100644 --- a/packages/ui/package.json +++ b/packages/ui/package.json @@ -73,6 +73,7 @@ "rehype-katex": "^7.0.1", "remark-math": "^6.0.0", "remend": "^1.2.1", + "shiki": "^3.23.0", "simple-git": "^3.28.0", "sonner": "^2.0.7", "strip-json-comments": "^5.0.3", diff --git a/packages/ui/src/components/chat/markdown/markdown-shiki.worker.ts b/packages/ui/src/components/chat/markdown/markdown-shiki.worker.ts new file mode 100644 index 00000000..c1012556 --- /dev/null +++ b/packages/ui/src/components/chat/markdown/markdown-shiki.worker.ts @@ -0,0 +1,55 @@ +/// + +import { bundledLanguages, createHighlighter, type BundledLanguage } from 'shiki'; +import { MARKDOWN_SHIKI_THEME, MARKDOWN_SHIKI_THEME_DEFINITION } from './markdownShikiThemeDefinition'; +import type { MarkdownWorkerRequest, MarkdownWorkerResponse } from './markdown-worker-protocol'; + +// Single shared highlighter for the worker. Languages load lazily on demand. +let highlighter: ReturnType | undefined; + +// Serialize work so language loading / tokenization never overlaps. +let queue = Promise.resolve(); + +const ensureHighlighter = (): ReturnType => { + highlighter ??= createHighlighter({ + // Cast: the theme is a CSS-variable TextMate theme; Shiki accepts the shape. + themes: [MARKDOWN_SHIKI_THEME_DEFINITION as unknown as Parameters[0]['themes'][number]], + langs: [], + }); + return highlighter; +}; + +self.onmessage = (event: MessageEvent) => { + const request = event.data; + if (request.type === 'init') { + ensureHighlighter(); + return; + } + queue = queue.then(() => highlight(request)).catch(() => {}); +}; + +async function highlight(request: Extract): Promise { + try { + const instance = await ensureHighlighter(); + let lang = request.lang in bundledLanguages ? request.lang : 'text'; + if (lang !== 'text' && !instance.getLoadedLanguages().includes(lang)) { + try { + await instance.loadLanguage(bundledLanguages[lang as BundledLanguage]); + } catch { + lang = 'text'; + } + } + const html = instance.codeToHtml(request.code, { + lang, + theme: MARKDOWN_SHIKI_THEME, + tabindex: false, + }); + post({ type: 'highlight', id: request.id, html }); + } catch (error) { + post({ type: 'error', id: request.id, message: error instanceof Error ? error.message : String(error) }); + } +} + +function post(response: MarkdownWorkerResponse): void { + self.postMessage(response); +} diff --git a/packages/ui/src/components/chat/markdown/markdown-worker-protocol.ts b/packages/ui/src/components/chat/markdown/markdown-worker-protocol.ts new file mode 100644 index 00000000..ec9230cb --- /dev/null +++ b/packages/ui/src/components/chat/markdown/markdown-worker-protocol.ts @@ -0,0 +1,13 @@ +// Message protocol for the markdown Shiki Web Worker. +// +// The worker tokenizes a complete code block off the main thread and returns +// ready-to-splice Shiki HTML. The theme is dependency-free and imported inside +// the worker directly, so it is not sent over postMessage. + +export type MarkdownWorkerRequest = + | { type: 'init' } + | { type: 'highlight'; id: number; code: string; lang: string }; + +export type MarkdownWorkerResponse = + | { type: 'highlight'; id: number; html: string } + | { type: 'error'; id: number; message: string }; diff --git a/packages/ui/src/components/chat/markdown/markdown-worker.ts b/packages/ui/src/components/chat/markdown/markdown-worker.ts new file mode 100644 index 00000000..3ed9b872 --- /dev/null +++ b/packages/ui/src/components/chat/markdown/markdown-worker.ts @@ -0,0 +1,53 @@ +import MarkdownShikiWorkerUrl from './markdown-shiki.worker.ts?worker&url'; +import type { MarkdownWorkerRequest, MarkdownWorkerResponse } from './markdown-worker-protocol'; + +// Main-thread client for the markdown Shiki 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. + +let worker: Worker | undefined; +let nextId = 0; +const pending = new Map void>(); + +const failAll = (): void => { + pending.forEach((resolve) => resolve(null)); + pending.clear(); + worker?.terminate(); + worker = undefined; +}; + +const getWorker = (): Worker | undefined => { + if (worker) return worker; + if (typeof window === 'undefined' || typeof Worker === 'undefined') return undefined; + try { + worker = new Worker(MarkdownShikiWorkerUrl, { type: 'module' }); + } catch { + return undefined; + } + worker.onmessage = (event: MessageEvent) => { + const resolve = pending.get(event.data.id); + if (!resolve) return; + pending.delete(event.data.id); + resolve(event.data.type === 'highlight' ? event.data.html : null); + }; + worker.onerror = failAll; + worker.onmessageerror = failAll; + worker.postMessage({ type: 'init' } satisfies MarkdownWorkerRequest); + return worker; +}; + +/** + * Highlight a complete code block in the worker. Resolves to Shiki `
` HTML,
+ * or `null` if highlighting is unavailable or failed (caller keeps plain code).
+ */
+export const highlightCodeInWorker = (code: string, lang: string): Promise => {
+  const instance = getWorker();
+  if (!instance) return Promise.resolve(null);
+  const id = ++nextId;
+  return new Promise((resolve) => {
+    pending.set(id, resolve);
+    instance.postMessage({ type: 'highlight', id, code, lang } satisfies MarkdownWorkerRequest);
+  });
+};
diff --git a/packages/ui/src/components/chat/markdown/markdownCore.ts b/packages/ui/src/components/chat/markdown/markdownCore.ts
index 502a9abd..1d324f8f 100644
--- a/packages/ui/src/components/chat/markdown/markdownCore.ts
+++ b/packages/ui/src/components/chat/markdown/markdownCore.ts
@@ -2,14 +2,9 @@ import { marked, type Tokens } from 'marked';
 import remend from 'remend';
 import katex from 'katex';
 import DOMPurify from 'dompurify';
-import {
-  getSharedHighlighter,
-  type DiffsThemeNames,
-  type SupportedLanguages,
-} from '@pierre/diffs';
 import { buildAgentMentionUrl, parseAgentHref, parseSkillHref } from '@/lib/messages/inlineMessageLinks';
 import { isVSCodeRuntime } from '@/lib/desktop';
-import { ensureMarkdownShikiTheme, MARKDOWN_SHIKI_THEME } from './markdownTheme';
+import { highlightCodeInWorker } from './markdown-worker';
 
 const escapeAttr = (value: string): string =>
   value.replace(/&/g, '&').replace(/"/g, '"').replace(//g, '>');
@@ -195,13 +190,6 @@ const highlightCodeBlocks = async (html: string): Promise => {
   const matches = [...html.matchAll(CODE_BLOCK_RE)];
   if (matches.length === 0) return html;
 
-  ensureMarkdownShikiTheme();
-  const highlighter = await getSharedHighlighter({
-    themes: [MARKDOWN_SHIKI_THEME as DiffsThemeNames],
-    langs: [],
-    preferredHighlighter: 'shiki-wasm',
-  });
-
   const lineLimit = isVSCodeRuntime() ? VSCODE_CODE_HIGHLIGHT_LINE_LIMIT : CODE_HIGHLIGHT_LINE_LIMIT;
 
   let result = html;
@@ -220,28 +208,13 @@ const highlightCodeBlocks = async (html: string): Promise => {
       continue;
     }
 
-    let lang = requested;
-
-    if (lang !== 'text' && !highlighter.getLoadedLanguages().includes(lang)) {
-      try {
-        await highlighter.loadLanguage(lang as SupportedLanguages);
-      } catch {
-        lang = 'text';
-      }
-    }
-
-    try {
-      const highlighted = highlighter
-        .codeToHtml(code, {
-          lang: lang as SupportedLanguages,
-          theme: MARKDOWN_SHIKI_THEME as DiffsThemeNames,
-          tabindex: false,
-        })
-        // Stamp the language so the decorate pass can show a header label.
-        .replace(/^
 highlighted);
-    } catch {
-      // Leave the original (escaped, sanitized) 
 in place.
+    // Tokenize off the main thread. On failure the worker resolves to null and
+    // we keep the original escaped 
 (no main-thread highlight).
+    const highlighted = await highlightCodeInWorker(code, requested);
+    if (highlighted) {
+      // Stamp the language so the decorate pass can show a header label.
+      const stamped = highlighted.replace(/^
 stamped);
     }
   }
 
diff --git a/packages/ui/src/components/chat/markdown/markdownShikiThemeDefinition.ts b/packages/ui/src/components/chat/markdown/markdownShikiThemeDefinition.ts
new file mode 100644
index 00000000..98ae1301
--- /dev/null
+++ b/packages/ui/src/components/chat/markdown/markdownShikiThemeDefinition.ts
@@ -0,0 +1,100 @@
+// Static, CSS-variable-driven Shiki theme definition.
+//
+// Token colors reference CSS custom properties (`--md-syntax-*`) instead of
+// concrete colors, so a highlighted code block does NOT need to be re-tokenized
+// when the app theme changes — only the CSS variables on the markdown container
+// update, and the browser repaints. This mirrors OpenCode's `var(--syntax-*)`
+// theme approach and keeps highlighting results cacheable across theme switches.
+//
+// This module is intentionally dependency-free (no `@pierre/diffs`, no React) so
+// it can be imported from inside the Shiki Web Worker bundle without dragging in
+// main-thread-only modules.
+
+export const MARKDOWN_SHIKI_THEME = 'openchamber-md';
+
+// Loosely typed on purpose: consumers (`@pierre/diffs` registration and the raw
+// Shiki worker) each cast to their own theme type. The shape is a standard
+// TextMate-style theme registration.
+export const MARKDOWN_SHIKI_THEME_DEFINITION = {
+  name: MARKDOWN_SHIKI_THEME,
+  colors: {
+    'editor.background': 'transparent',
+    'editor.foreground': 'var(--md-syntax-foreground)',
+  },
+  tokenColors: [
+    {
+      scope: ['comment', 'punctuation.definition.comment', 'string.comment'],
+      settings: { foreground: 'var(--md-syntax-comment)', fontStyle: 'italic' },
+    },
+    {
+      scope: ['string', 'punctuation.definition.string', 'string.template'],
+      settings: { foreground: 'var(--md-syntax-string)' },
+    },
+    {
+      scope: ['constant.numeric', 'constant.language', 'constant.character', 'constant'],
+      settings: { foreground: 'var(--md-syntax-number)' },
+    },
+    {
+      scope: ['keyword', 'storage', 'storage.type', 'storage.modifier', 'keyword.control'],
+      settings: { foreground: 'var(--md-syntax-keyword)' },
+    },
+    {
+      scope: ['keyword.operator', 'punctuation.separator', 'punctuation.terminator'],
+      settings: { foreground: 'var(--md-syntax-operator)' },
+    },
+    {
+      scope: ['entity.name.function', 'support.function', 'meta.function-call'],
+      settings: { foreground: 'var(--md-syntax-function)' },
+    },
+    {
+      scope: [
+        'entity.name.type',
+        'entity.name.class',
+        'support.type',
+        'support.class',
+        'entity.other.inherited-class',
+      ],
+      settings: { foreground: 'var(--md-syntax-type)' },
+    },
+    {
+      scope: ['variable', 'variable.other', 'variable.parameter', 'meta.definition.variable'],
+      settings: { foreground: 'var(--md-syntax-variable)' },
+    },
+    {
+      scope: ['variable.other.property', 'meta.property-name', 'support.type.property-name'],
+      settings: { foreground: 'var(--md-syntax-property)' },
+    },
+    {
+      scope: ['entity.name.tag', 'punctuation.definition.tag'],
+      settings: { foreground: 'var(--md-syntax-keyword)' },
+    },
+    {
+      scope: ['entity.other.attribute-name'],
+      settings: { foreground: 'var(--md-syntax-property)' },
+    },
+    {
+      scope: ['markup.bold', 'punctuation.definition.bold'],
+      settings: { fontStyle: 'bold' },
+    },
+    {
+      scope: ['markup.italic', 'punctuation.definition.italic'],
+      settings: { fontStyle: 'italic' },
+    },
+    {
+      scope: ['markup.heading', 'markup.heading entity.name'],
+      settings: { foreground: 'var(--md-syntax-keyword)', fontStyle: 'bold' },
+    },
+    {
+      scope: ['markup.inserted', 'punctuation.definition.inserted'],
+      settings: { foreground: 'var(--md-syntax-inserted)' },
+    },
+    {
+      scope: ['markup.deleted', 'punctuation.definition.deleted'],
+      settings: { foreground: 'var(--md-syntax-deleted)' },
+    },
+    {
+      scope: ['invalid', 'invalid.illegal'],
+      settings: { foreground: 'var(--md-syntax-deleted)' },
+    },
+  ],
+} as const;
diff --git a/packages/ui/src/components/chat/markdown/markdownTheme.ts b/packages/ui/src/components/chat/markdown/markdownTheme.ts
index fbdd370a..f78161cd 100644
--- a/packages/ui/src/components/chat/markdown/markdownTheme.ts
+++ b/packages/ui/src/components/chat/markdown/markdownTheme.ts
@@ -1,108 +1,30 @@
 import { registerCustomTheme, type ThemeRegistrationResolved } from '@pierre/diffs';
 import type { Theme } from '@/types/theme';
+import { MARKDOWN_SHIKI_THEME, MARKDOWN_SHIKI_THEME_DEFINITION } from './markdownShikiThemeDefinition';
 
-// Name of the static Shiki theme we register once. Its token colors reference
-// CSS custom properties (`--md-syntax-*`) instead of concrete colors, so a
-// highlighted code block does NOT need to be re-tokenized when the app theme
-// changes — only the CSS variables on the markdown container update, and the
-// browser repaints. This mirrors OpenCode's `var(--syntax-*)` theme approach
-// and keeps highlighting results cacheable across theme switches.
-export const MARKDOWN_SHIKI_THEME = 'openchamber-md';
+// The static Shiki theme name. Its definition (token colors referencing
+// `--md-syntax-*` CSS variables) lives in the dependency-free
+// `markdownShikiThemeDefinition` module so it can also be imported inside the
+// Shiki Web Worker. See that module for the rationale.
+export { MARKDOWN_SHIKI_THEME };
 
 let registered = false;
 
 /**
- * Register the static, CSS-variable-driven Shiki theme. Safe to call multiple
- * times; only the first call registers.
+ * Register the static, CSS-variable-driven Shiki theme with `@pierre/diffs`.
+ * Safe to call multiple times; only the first call registers.
+ *
+ * NOTE: markdown code highlighting now runs through the dedicated Shiki worker
+ * (`markdown-worker`), which uses the raw theme definition directly. This
+ * registration remains only for any `@pierre/diffs`-based consumer of the
+ * `openchamber-md` theme name.
  */
 export const ensureMarkdownShikiTheme = (): void => {
   if (registered) return;
   registered = true;
 
   registerCustomTheme(MARKDOWN_SHIKI_THEME, () =>
-    Promise.resolve({
-      name: MARKDOWN_SHIKI_THEME,
-      colors: {
-        'editor.background': 'transparent',
-        'editor.foreground': 'var(--md-syntax-foreground)',
-      },
-      tokenColors: [
-        {
-          scope: ['comment', 'punctuation.definition.comment', 'string.comment'],
-          settings: { foreground: 'var(--md-syntax-comment)', fontStyle: 'italic' },
-        },
-        {
-          scope: ['string', 'punctuation.definition.string', 'string.template'],
-          settings: { foreground: 'var(--md-syntax-string)' },
-        },
-        {
-          scope: ['constant.numeric', 'constant.language', 'constant.character', 'constant'],
-          settings: { foreground: 'var(--md-syntax-number)' },
-        },
-        {
-          scope: ['keyword', 'storage', 'storage.type', 'storage.modifier', 'keyword.control'],
-          settings: { foreground: 'var(--md-syntax-keyword)' },
-        },
-        {
-          scope: ['keyword.operator', 'punctuation.separator', 'punctuation.terminator'],
-          settings: { foreground: 'var(--md-syntax-operator)' },
-        },
-        {
-          scope: ['entity.name.function', 'support.function', 'meta.function-call'],
-          settings: { foreground: 'var(--md-syntax-function)' },
-        },
-        {
-          scope: [
-            'entity.name.type',
-            'entity.name.class',
-            'support.type',
-            'support.class',
-            'entity.other.inherited-class',
-          ],
-          settings: { foreground: 'var(--md-syntax-type)' },
-        },
-        {
-          scope: ['variable', 'variable.other', 'variable.parameter', 'meta.definition.variable'],
-          settings: { foreground: 'var(--md-syntax-variable)' },
-        },
-        {
-          scope: ['variable.other.property', 'meta.property-name', 'support.type.property-name'],
-          settings: { foreground: 'var(--md-syntax-property)' },
-        },
-        {
-          scope: ['entity.name.tag', 'punctuation.definition.tag'],
-          settings: { foreground: 'var(--md-syntax-keyword)' },
-        },
-        {
-          scope: ['entity.other.attribute-name'],
-          settings: { foreground: 'var(--md-syntax-property)' },
-        },
-        {
-          scope: ['markup.bold', 'punctuation.definition.bold'],
-          settings: { fontStyle: 'bold' },
-        },
-        {
-          scope: ['markup.italic', 'punctuation.definition.italic'],
-          settings: { fontStyle: 'italic' },
-        },
-        {
-          scope: ['markup.heading', 'markup.heading entity.name'],
-          settings: { foreground: 'var(--md-syntax-keyword)', fontStyle: 'bold' },
-        },
-        {
-          scope: ['markup.inserted', 'punctuation.definition.inserted'],
-          settings: { foreground: 'var(--md-syntax-inserted)' },
-        },
-        {
-          scope: ['markup.deleted', 'punctuation.definition.deleted'],
-          settings: { foreground: 'var(--md-syntax-deleted)' },
-        },
-        {
-          scope: ['invalid', 'invalid.illegal'],
-          settings: { foreground: 'var(--md-syntax-deleted)' },
-        },
-      ],
-    } as unknown as ThemeRegistrationResolved),
+    Promise.resolve(MARKDOWN_SHIKI_THEME_DEFINITION as unknown as ThemeRegistrationResolved),
   );
 };