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;