fix: reduce code block highlight flicker

Preserve layout with invisible placeholders while worker highlighting loads
Reuse cached highlighted lines on first render when available
Differentiate loading from failed highlighting results for better fallback behavior

packages/ui/src/components/chat/DiffPreview.tsx

@@ -3,17 +3,29 @@ import { cn } from '@/lib/utils';
 import { getLanguageFromExtension } from '@/lib/toolHelpers';
 import { useThemeSystem } from '@/contexts/useThemeSystem';
 import { getMarkdownSyntaxVars } from '@/components/chat/markdown/markdownSyntaxVars';
-import { useWorkerHighlightedLines } from '@/components/code/useWorkerHighlightedLines';
+import {
+  useWorkerHighlightedLines,
+  type WorkerHighlightedLinesResult,
+} from '@/components/code/useWorkerHighlightedLines';
 import { parseDiffToUnified } from './message/toolRenderers';
 
-// One highlighted line: swaps in worker-tokenized inner HTML when ready, falls
-// back to plain text while loading or on failure.
-const CodeLineContent: React.FC<{ content: string; html: string | undefined }> = ({ content, html }) =>
-  html !== undefined ? (
-    <span className="whitespace-pre-wrap break-all" dangerouslySetInnerHTML={{ __html: html }} />
-  ) : (
-    <span className="whitespace-pre-wrap break-all">{content}</span>
-  );
+// Keep the line's layout stable while a cold worker request finishes. Plain
+// text appears only if highlighting fails, avoiding a visible color flash.
+interface CodeLineContentProps {
+  content: string;
+  html: string | undefined;
+  status: WorkerHighlightedLinesResult['status'];
+}
+
+const CodeLineContent: React.FC<CodeLineContentProps> = ({ content, html, status }) => {
+  if (status === 'ready' && html !== undefined) {
+    return <span className="whitespace-pre-wrap break-all" dangerouslySetInnerHTML={{ __html: html }} />;
+  }
+  if (status === 'loading') {
+    return <span aria-hidden className="invisible whitespace-pre-wrap break-all">{content}</span>;
+  }
+  return <span className="whitespace-pre-wrap break-all">{content}</span>;
+};
 
 interface DiffPreviewProps {
     diff: string;
@@ -44,7 +56,7 @@ export const DiffPreview: React.FC<DiffPreviewProps> = ({ diff, filePath }) => {
 
                     <div>
                         {hunk.lines.map((line, lineIdx) => {
-                            const html = highlighted?.[lineCursor];
+                            const html = highlighted.lines?.[lineCursor];
                             lineCursor += 1;
                             return (
                                 <div
@@ -67,7 +79,7 @@ export const DiffPreview: React.FC<DiffPreviewProps> = ({ diff, filePath }) => {
                                         {line.lineNumber || ''}
                                     </span>
                                     <div className="flex-1 min-w-0">
-                                        <CodeLineContent content={line.content} html={html} />
+                                        <CodeLineContent content={line.content} html={html} status={highlighted.status} />
                                     </div>
                                 </div>
                             );
@@ -106,7 +118,11 @@ export const WritePreview: React.FC<WritePreviewProps> = ({ content, filePath })
                             {lineIdx + 1}
                         </span>
                         <div className="flex-1 min-w-0">
-                            <CodeLineContent content={line || ' '} html={highlighted?.[lineIdx]} />
+                            <CodeLineContent
+                                content={line || ' '}
+                                html={highlighted.lines?.[lineIdx]}
+                                status={highlighted.status}
+                            />
                         </div>
                     </div>
                 ))}

packages/ui/src/components/chat/markdown/markdown-worker.ts

@@ -167,6 +167,12 @@ export const highlightLinesInWorker = async (code: string, lang: string): Promis
   return result?.type === 'highlightLines' ? result.lines : null;
 };
 
+/** Return an already-tokenized line result without scheduling a worker request. */
+export const getCachedHighlightedLines = (code: string, lang: string): string[] | null => {
+  const cached = resultCache.get(cacheKeyFor('highlightLines', lang, code));
+  return cached?.type === 'highlightLines' ? cached.lines : null;
+};
+
 /**
  * Tokenize `code` with the given resolved TextMate theme and return per-line
  * styled runs with offsets — for building CodeMirror decorations that match the

packages/ui/src/components/chat/message/parts/VirtualizedCodeBlock.tsx

@@ -54,7 +54,8 @@ export const VirtualizedCodeBlock: React.FC<VirtualizedCodeBlockProps> = React.m
   const syntaxVars = React.useMemo(() => getMarkdownSyntaxVars(currentTheme), [currentTheme]);
   // Tokenize the whole block in one worker call; rows index into the result.
   const fullText = React.useMemo(() => lines.map((line) => line.text).join('\n'), [lines]);
-  const highlighted = useWorkerHighlightedLines(fullText, language);
+  const highlightResult = useWorkerHighlightedLines(fullText, language);
+  const highlighted = highlightResult.lines;
 
   const shouldVirtualize = lines.length > VIRTUALIZE_THRESHOLD;
 

packages/ui/src/components/code/useWorkerHighlightedLines.ts

@@ -1,26 +1,60 @@
 import React from 'react';
-import { highlightLinesInWorker } from '@/components/chat/markdown/markdown-worker';
+import {
+  getCachedHighlightedLines,
+  highlightLinesInWorker,
+} from '@/components/chat/markdown/markdown-worker';
 
 // Tokenize a whole block ONCE in the Shiki worker and expose per-line inner
 // HTML. For per-line layouts (diffs, gutters, virtualization) that would
-// otherwise spawn one highlighter per row. Returns `null` until the first
-// result lands (or permanently on failure) — callers render plain text then.
+// otherwise spawn one highlighter per row. Cached results are available on the
+// first render; cold requests distinguish loading from permanent failure so
+// callers can choose whether to reveal their plain-text fallback.
 //
 // Whole-block tokenization also restores cross-line syntax context (multi-line
 // strings / comments) that independent per-line highlighting loses.
-export const useWorkerHighlightedLines = (code: string, language: string): string[] | null => {
-  const [lines, setLines] = React.useState<string[] | null>(null);
+export type WorkerHighlightedLinesResult =
+  | { status: 'loading'; lines: null }
+  | { status: 'ready'; lines: string[] }
+  | { status: 'failed'; lines: null };
+
+type HighlightState = WorkerHighlightedLinesResult & {
+  code: string;
+  language: string;
+};
+
+const getHighlightState = (code: string, language: string): HighlightState => {
+  const lines = getCachedHighlightedLines(code, language);
+  return lines
+    ? { status: 'ready', lines, code, language }
+    : { status: 'loading', lines: null, code, language };
+};
+
+export const useWorkerHighlightedLines = (code: string, language: string): WorkerHighlightedLinesResult => {
+  const normalizedLanguage = (language || 'text').toLowerCase();
+  const [state, setState] = React.useState<HighlightState>(() => getHighlightState(code, normalizedLanguage));
 
   React.useEffect(() => {
+    const cached = getCachedHighlightedLines(code, normalizedLanguage);
+    if (cached) {
+      setState({ status: 'ready', lines: cached, code, language: normalizedLanguage });
+      return;
+    }
+
+    setState({ status: 'loading', lines: null, code, language: normalizedLanguage });
     let active = true;
-    setLines(null);
-    void highlightLinesInWorker(code, (language || 'text').toLowerCase()).then((result) => {
-      if (active) setLines(result);
+    void highlightLinesInWorker(code, normalizedLanguage).then((lines) => {
+      if (!active) return;
+      setState(lines
+        ? { status: 'ready', lines, code, language: normalizedLanguage }
+        : { status: 'failed', lines: null, code, language: normalizedLanguage });
     });
     return () => {
       active = false;
     };
-  }, [code, language]);
+  }, [code, normalizedLanguage]);
 
-  return lines;
+  if (state.code !== code || state.language !== normalizedLanguage) {
+    return getHighlightState(code, normalizedLanguage);
+  }
+  return state;
 };