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/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;
 };