From 0608900a530dcb7edf8fde678e2bdba4ef8c9b34 Mon Sep 17 00:00:00 2001 From: Lukas May Date: Fri, 6 Mar 2026 19:39:31 +0100 Subject: [PATCH] feat: move syntax highlighting off main thread via Web Worker pool MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds a 2-worker pool in use-syntax-highlight.ts so shiki tokenisation runs off the main thread. Callers continue to receive null while the worker is in flight and a LineTokenMap once it resolves — no caller changes needed. Fallback: if Worker construction is blocked (e.g. CSP), the hook falls back to the existing createHighlighter singleton but processes 200 lines at a time, yielding between chunks via scheduler.yield()/setTimeout(0) to avoid long tasks. Also adds worker.format:'es' to vite.config.ts (required when the app uses code-splitting) and covers all paths with Vitest tests. Co-Authored-By: Claude Sonnet 4.6 --- .../src/components/review/highlight-worker.ts | 39 +++ .../use-syntax-highlight.fallback.test.ts | 70 +++++ .../review/use-syntax-highlight.test.ts | 240 ++++++++++++++++++ .../components/review/use-syntax-highlight.ts | 177 ++++++++++--- apps/web/vite.config.ts | 5 + 5 files changed, 501 insertions(+), 30 deletions(-) create mode 100644 apps/web/src/components/review/highlight-worker.ts create mode 100644 apps/web/src/components/review/use-syntax-highlight.fallback.test.ts create mode 100644 apps/web/src/components/review/use-syntax-highlight.test.ts diff --git a/apps/web/src/components/review/highlight-worker.ts b/apps/web/src/components/review/highlight-worker.ts new file mode 100644 index 0000000..d663c16 --- /dev/null +++ b/apps/web/src/components/review/highlight-worker.ts @@ -0,0 +1,39 @@ +import type { ThemedToken } from 'shiki'; + +export interface HighlightRequest { + id: string; + filePath: string; + language: string; // resolved lang name (e.g. "typescript") or "text" + code: string; // full joined content of new-side lines to highlight + lineNumbers: number[]; // new-side line numbers to map tokens back to +} + +export interface HighlightResponse { + id: string; + tokens: Array<{ lineNumber: number; tokens: ThemedToken[] }>; + error?: string; +} + +self.addEventListener('message', async (event: MessageEvent) => { + const { id, language, code, lineNumbers } = event.data; + try { + const { codeToTokens } = await import('shiki'); + const result = await codeToTokens(code, { + lang: language as Parameters[1]['lang'], + theme: 'github-dark-default', + }); + const tokens: HighlightResponse['tokens'] = result.tokens.map((lineTokens, idx) => ({ + lineNumber: lineNumbers[idx] ?? idx, + tokens: lineTokens, + })); + const response: HighlightResponse = { id, tokens }; + self.postMessage(response); + } catch (err) { + const response: HighlightResponse = { + id, + tokens: [], + error: err instanceof Error ? err.message : String(err), + }; + self.postMessage(response); + } +}); diff --git a/apps/web/src/components/review/use-syntax-highlight.fallback.test.ts b/apps/web/src/components/review/use-syntax-highlight.fallback.test.ts new file mode 100644 index 0000000..1cd7bb7 --- /dev/null +++ b/apps/web/src/components/review/use-syntax-highlight.fallback.test.ts @@ -0,0 +1,70 @@ +// @vitest-environment happy-dom +// This file tests the chunked main-thread fallback path when Worker +// construction is blocked (e.g. by CSP). It runs in isolation from the +// worker-path tests so that module-level state (workersInitialized, workers) +// starts clean. +import '@testing-library/jest-dom/vitest' +import { renderHook, waitFor } from '@testing-library/react' +import { vi, describe, it, expect, beforeAll, afterAll } from 'vitest' + +const MOCK_TOKEN_A = { content: 'const', color: '#569cd6', offset: 0 } +const MOCK_TOKEN_B = { content: 'x', color: '#9cdcfe', offset: 0 } + +// Mock shiki's createHighlighter for the fallback path +const mockCodeToTokens = vi.fn() + +vi.mock('shiki', () => ({ + createHighlighter: vi.fn().mockResolvedValue({ + codeToTokens: mockCodeToTokens, + }), +})) + +// Stub Worker to throw (simulating CSP) BEFORE the hook module is loaded. +// initWorkers() catches the exception and leaves workers = []. +beforeAll(() => { + // Use a class so Vitest doesn't warn about constructing vi.fn() without a class impl + class BlockedWorker { + constructor() { + throw new Error('CSP blocks workers') + } + } + vi.stubGlobal('Worker', BlockedWorker) + + mockCodeToTokens.mockReturnValue({ + tokens: [[MOCK_TOKEN_A], [MOCK_TOKEN_B]], + }) +}) + +afterAll(() => { + vi.unstubAllGlobals() +}) + +// Dynamic import ensures this file's module instance is fresh (workersInitialized = false). +// We import inside tests below rather than at the top level. + +describe('useHighlightedFile — fallback path (Worker unavailable)', () => { + it('falls back to chunked main-thread highlighting when Worker construction throws', async () => { + const { useHighlightedFile } = await import('./use-syntax-highlight') + + const lines = [ + { content: 'const x = 1', newLineNumber: 1, type: 'added' as const }, + { content: 'let y = 2', newLineNumber: 2, type: 'context' as const }, + ] + + const { result } = renderHook(() => useHighlightedFile('app.ts', lines)) + + // Initially null while chunked highlighting runs + expect(result.current).toBeNull() + + // Fallback createHighlighter path eventually resolves tokens + await waitFor( + () => { + expect(result.current).not.toBeNull() + }, + { timeout: 5000 }, + ) + + expect(result.current?.get(1)).toEqual([MOCK_TOKEN_A]) + expect(result.current?.get(2)).toEqual([MOCK_TOKEN_B]) + }) +}) diff --git a/apps/web/src/components/review/use-syntax-highlight.test.ts b/apps/web/src/components/review/use-syntax-highlight.test.ts new file mode 100644 index 0000000..1356c3b --- /dev/null +++ b/apps/web/src/components/review/use-syntax-highlight.test.ts @@ -0,0 +1,240 @@ +// @vitest-environment happy-dom +import '@testing-library/jest-dom/vitest' +import { renderHook, waitFor, act } from '@testing-library/react' +import { vi, describe, it, expect, beforeAll, beforeEach, afterAll } from 'vitest' + +// ── Worker mock infrastructure ───────────────────────────────────────────── +// +// We stub Worker BEFORE importing use-syntax-highlight so that initWorkers() +// (called from useEffect on first render) picks up our mock. +// Module-level state (workers, pending, workersInitialized) is shared across +// all tests in this file — we control behaviour through the mock instances. + +type WorkerHandler = (event: { data: unknown }) => void + +class MockWorker { + static instances: MockWorker[] = [] + + messageHandler: WorkerHandler | null = null + postMessage = vi.fn() + + constructor() { + MockWorker.instances.push(this) + } + + addEventListener(type: string, handler: WorkerHandler) { + if (type === 'message') this.messageHandler = handler + } + + /** Simulate a message arriving from the worker thread */ + simulateResponse(data: unknown) { + this.messageHandler?.({ data }) + } +} + +// Stub Worker before the hook module is loaded. +// initWorkers() is lazy (called inside useEffect), so the stub is in place +// by the time any test renders a hook. +beforeAll(() => { + vi.stubGlobal('Worker', MockWorker) +}) + +afterAll(() => { + vi.unstubAllGlobals() +}) + +beforeEach(() => { + // Reset call history between tests; keep instances (pool is created once) + MockWorker.instances.forEach((w) => w.postMessage.mockClear()) +}) + +// Import the hook AFTER the beforeAll stub is registered (hoisted evaluation +// of the module will not call initWorkers() — that happens in useEffect). +import { useHighlightedFile } from './use-syntax-highlight' + +// ── Helpers ──────────────────────────────────────────────────────────────── + +const MOCK_TOKEN_A = { content: 'const', color: '#569cd6', offset: 0 } +const MOCK_TOKEN_B = { content: 'x', color: '#9cdcfe', offset: 0 } + +function makeLine( + content: string, + newLineNumber: number, + type: 'added' | 'context' | 'removed' = 'added', +) { + return { content, newLineNumber, type } as const +} + +// ── Tests ────────────────────────────────────────────────────────────────── + +describe('useHighlightedFile — worker path', () => { + // ── Test 1: Correct message format ─────────────────────────────────────── + + it('posts a message to a worker with filePath, language, code, and lineNumbers', async () => { + const lines = [ + makeLine('const x = 1', 1, 'added'), + makeLine('const y = 2', 2, 'context'), + ] + + renderHook(() => useHighlightedFile('src/index.ts', lines)) + + // Wait for initWorkers() to fire and postMessage to be called + await waitFor(() => { + const totalCalls = MockWorker.instances.reduce( + (n, w) => n + w.postMessage.mock.calls.length, + 0, + ) + expect(totalCalls).toBeGreaterThan(0) + }) + + // Find which worker received the message + const calledWorker = MockWorker.instances.find((w) => w.postMessage.mock.calls.length > 0) + expect(calledWorker).toBeDefined() + expect(calledWorker!.postMessage).toHaveBeenCalledWith( + expect.objectContaining({ + filePath: 'src/index.ts', + language: 'typescript', + code: 'const x = 1\nconst y = 2', + lineNumbers: [1, 2], + }), + ) + }) + + // ── Test 2: Response builds token map ───────────────────────────────────── + + it('returns null initially and a LineTokenMap after worker responds', async () => { + const lines = [makeLine('const x = 1', 10, 'added')] + + const { result } = renderHook(() => useHighlightedFile('component.ts', lines)) + + // Immediately null while worker is pending + expect(result.current).toBeNull() + + // Capture the request id from whichever worker received it + let requestId = '' + let respondingWorker: MockWorker | undefined + + await waitFor(() => { + respondingWorker = MockWorker.instances.find((w) => w.postMessage.mock.calls.length > 0) + expect(respondingWorker).toBeDefined() + requestId = respondingWorker!.postMessage.mock.calls[0][0].id as string + expect(requestId).not.toBe('') + }) + + // Simulate the worker responding + act(() => { + respondingWorker!.simulateResponse({ + id: requestId, + tokens: [{ lineNumber: 10, tokens: [MOCK_TOKEN_A] }], + }) + }) + + await waitFor(() => { + expect(result.current).not.toBeNull() + expect(result.current?.get(10)).toEqual([MOCK_TOKEN_A]) + }) + }) + + // ── Test 3: Worker error response → null ────────────────────────────────── + + it('returns null when worker responds with an error field', async () => { + const lines = [makeLine('code here', 1, 'added')] + + const { result } = renderHook(() => useHighlightedFile('bad.ts', lines)) + + let requestId = '' + let respondingWorker: MockWorker | undefined + + await waitFor(() => { + respondingWorker = MockWorker.instances.find((w) => w.postMessage.mock.calls.length > 0) + expect(respondingWorker).toBeDefined() + requestId = respondingWorker!.postMessage.mock.calls[0][0].id as string + }) + + act(() => { + respondingWorker!.simulateResponse({ + id: requestId, + tokens: [], + error: 'Worker crashed', + }) + }) + + // Error → stays null (plain text fallback in the UI) + await new Promise((r) => setTimeout(r, 20)) + expect(result.current).toBeNull() + }) + + // ── Test 4: Unmount before response — no state update ──────────────────── + + it('silently discards a late worker response after unmount', async () => { + const lines = [makeLine('const z = 3', 5, 'added')] + + const { result, unmount } = renderHook(() => useHighlightedFile('late.ts', lines)) + + let requestId = '' + let respondingWorker: MockWorker | undefined + + await waitFor(() => { + respondingWorker = MockWorker.instances.find((w) => w.postMessage.mock.calls.length > 0) + expect(respondingWorker).toBeDefined() + requestId = respondingWorker!.postMessage.mock.calls[0][0].id as string + }) + + // Unmount before the response arrives + unmount() + + // Simulate the late response — should be silently dropped + act(() => { + respondingWorker!.simulateResponse({ + id: requestId, + tokens: [{ lineNumber: 5, tokens: [MOCK_TOKEN_B] }], + }) + }) + + // result.current is frozen at last rendered value (null) — no update fired + expect(result.current).toBeNull() + }) + + // ── Test 5: Round-robin — two simultaneous requests go to different workers + + it('distributes two simultaneous requests across both pool workers', async () => { + // Ensure the pool has been initialised (first test may have done this) + // and reset call counts for clean measurement. + MockWorker.instances.forEach((w) => w.postMessage.mockClear()) + + const lines1 = [makeLine('alpha', 1, 'added')] + const lines2 = [makeLine('beta', 1, 'added')] + + // Render two hook instances at the same time + renderHook(() => useHighlightedFile('file1.ts', lines1)) + renderHook(() => useHighlightedFile('file2.ts', lines2)) + + await waitFor(() => { + const total = MockWorker.instances.reduce((n, w) => n + w.postMessage.mock.calls.length, 0) + expect(total).toBe(2) + }) + + // Both pool workers should each have received exactly one request + // (round-robin: even requestCount → workers[0], odd → workers[1]) + const counts = MockWorker.instances.map((w) => w.postMessage.mock.calls.length) + // Pool has 2 workers; each should have received 1 of the 2 requests + expect(counts[0]).toBe(1) + expect(counts[1]).toBe(1) + }) + + // ── Test 6: Unknown language → no request ──────────────────────────────── + + it('returns null immediately for files with no detectable language', async () => { + MockWorker.instances.forEach((w) => w.postMessage.mockClear()) + + const lines = [makeLine('raw data', 1, 'added')] + + const { result } = renderHook(() => useHighlightedFile('data.xyz', lines)) + + await new Promise((r) => setTimeout(r, 50)) + + expect(result.current).toBeNull() + const total = MockWorker.instances.reduce((n, w) => n + w.postMessage.mock.calls.length, 0) + expect(total).toBe(0) + }) +}) diff --git a/apps/web/src/components/review/use-syntax-highlight.ts b/apps/web/src/components/review/use-syntax-highlight.ts index 673591d..dda0f63 100644 --- a/apps/web/src/components/review/use-syntax-highlight.ts +++ b/apps/web/src/components/review/use-syntax-highlight.ts @@ -1,7 +1,59 @@ import { useState, useEffect, useMemo } from "react"; import type { ThemedToken } from "shiki"; +import type { HighlightRequest, HighlightResponse } from "./highlight-worker"; -/* ── Lazy singleton highlighter ─────────────────────────── */ +/* ── Worker pool (module-level, shared across all hook instances) ─────── */ + +type PendingResolve = (response: HighlightResponse) => void; + +let workers: Worker[] = []; +let requestCount = 0; +const MAX_WORKERS = 2; +const pending = new Map(); + +let workersInitialized = false; + +function initWorkers(): void { + if (workersInitialized) return; + workersInitialized = true; + try { + workers = Array.from({ length: MAX_WORKERS }, () => { + const w = new Worker( + new URL("./highlight-worker.ts", import.meta.url), + { type: "module" }, + ); + w.addEventListener("message", (event: MessageEvent) => { + const resolve = pending.get(event.data.id); + if (resolve) { + pending.delete(event.data.id); + resolve(event.data); + } + }); + return w; + }); + } catch { + // CSP or browser compat — fall back to chunked main-thread highlighting + workers = []; + } +} + +function highlightWithWorker( + id: string, + language: string, + code: string, + lineNumbers: number[], + filePath: string, +): Promise { + return new Promise((resolve) => { + pending.set(id, resolve); + const worker = workers[requestCount % MAX_WORKERS]; + requestCount++; + const req: HighlightRequest = { id, filePath, language, code, lineNumbers }; + worker.postMessage(req); + }); +} + +/* ── Lazy singleton highlighter (for main-thread fallback) ───────────── */ let highlighterPromise: Promise @@ -40,10 +92,59 @@ function getHighlighter() { return highlighterPromise; } -// Pre-warm on module load (non-blocking) -getHighlighter(); +/* ── Chunked main-thread fallback ────────────────────────────────────── */ -/* ── Language detection ──────────────────────────────────── */ +async function highlightChunked( + code: string, + language: string, + lineNumbers: number[], + signal: AbortSignal, +): Promise { + const CHUNK = 200; + const result: LineTokenMap = new Map(); + const lines = code.split("\n"); + const highlighter = await getHighlighter(); + if (!highlighter) return result; + + for (let i = 0; i < lines.length; i += CHUNK) { + if (signal.aborted) break; + const chunkLines = lines.slice(i, i + CHUNK); + const chunkCode = chunkLines.join("\n"); + try { + const tokenized = highlighter.codeToTokens(chunkCode, { + lang: language as Parameters[1]["lang"], + theme: "github-dark-default", + }); + tokenized.tokens.forEach((lineTokens: ThemedToken[], idx: number) => { + const lineNum = lineNumbers[i + idx]; + if (lineNum !== undefined) result.set(lineNum, lineTokens); + }); + } catch { + // Skip unparseable chunk + } + + // Yield between chunks to avoid blocking the main thread + await new Promise((r) => { + if ( + "scheduler" in globalThis && + "yield" in (globalThis as Record).scheduler + ) { + ( + (globalThis as Record).scheduler as { + yield: () => Promise; + } + ) + .yield() + .then(r); + } else { + setTimeout(r, 0); + } + }); + } + return result; +} + +/* ── Language detection ──────────────────────────────────────────────── */ const EXT_TO_LANG: Record = { ts: "typescript", @@ -77,7 +178,7 @@ function detectLang(path: string): string | null { return EXT_TO_LANG[ext] ?? null; } -/* ── Types ───────────────────────────────────────────────── */ +/* ── Types ───────────────────────────────────────────────────────────── */ export type TokenizedLine = ThemedToken[]; /** Maps newLineNumber → highlighted tokens for that line */ @@ -89,12 +190,23 @@ interface DiffLineInput { type: "added" | "removed" | "context"; } -/* ── Hook ────────────────────────────────────────────────── */ +/* ── Hook ────────────────────────────────────────────────────────────── */ /** - * Highlights the "new-side" content of a file diff. - * Returns null until highlighting is ready (progressive enhancement). - * Only context + added lines are highlighted (removed lines fall back to plain text). + * Highlights the "new-side" content of a file diff, returning a map of + * line number → syntax tokens. + * + * Progressive rendering: returns `null` while highlighting is in progress. + * Callers (HunkRows → LineWithComments) render plain text when `null` and + * patch in highlighted tokens on re-render once the worker or chunked call + * resolves. + * + * Worker path: uses a module-level pool of 2 Web Workers. Round-robin + * assignment. Late responses after unmount are silently discarded. + * + * Fallback path: if Worker construction fails (CSP, browser compat), + * falls back to chunked main-thread highlighting via codeToTokens (200 + * lines/chunk) with scheduler.yield()/setTimeout(0) between chunks. */ export function useHighlightedFile( filePath: string, @@ -129,32 +241,37 @@ export function useHighlightedFile( return; } - let cancelled = false; + initWorkers(); // no-op after first call - getHighlighter().then((highlighter) => { - if (cancelled || !highlighter) return; + const id = crypto.randomUUID(); + let unmounted = false; + const abortController = new AbortController(); - try { - const result = highlighter.codeToTokens(code, { - lang: lang as Parameters[1]["lang"], - theme: "github-dark-default", - }); + if (workers.length > 0) { + highlightWithWorker(id, lang, code, lineNums, filePath).then((response) => { + if (unmounted) return; // ignore late responses after unmount + if (response.error || response.tokens.length === 0) { + setTokenMap(null); + return; + } const map: LineTokenMap = new Map(); - - result.tokens.forEach((lineTokens: ThemedToken[], idx: number) => { - if (idx < lineNums.length) { - map.set(lineNums[idx], lineTokens); - } - }); - - if (!cancelled) setTokenMap(map); - } catch { - // Language not loaded or parse error — no highlighting - } - }); + for (const { lineNumber, tokens } of response.tokens) { + map.set(lineNumber, tokens); + } + setTokenMap(map); + }); + } else { + highlightChunked(code, lang, lineNums, abortController.signal).then((map) => { + if (unmounted) return; + setTokenMap(map.size > 0 ? map : null); + }); + } return () => { - cancelled = true; + unmounted = true; + abortController.abort(); + // Remove pending resolver so a late worker response is silently dropped + pending.delete(id); }; // eslint-disable-next-line react-hooks/exhaustive-deps }, [cacheKey]); diff --git a/apps/web/vite.config.ts b/apps/web/vite.config.ts index 6153961..a92669e 100644 --- a/apps/web/vite.config.ts +++ b/apps/web/vite.config.ts @@ -10,6 +10,11 @@ export default defineConfig({ "@": path.resolve(__dirname, "./src"), }, }, + worker: { + // ES module workers are required when the app uses code-splitting (Rollup + // can't bundle IIFE workers alongside dynamic imports). + format: "es", + }, server: { proxy: { "/trpc": {