🤖 fix: stop code-block highlight flashing while streaming

src/browser/features/Messages/MarkdownComponents.test.tsx

@@ -3,7 +3,13 @@ import { afterEach, beforeEach, describe, expect, mock, test } from "bun:test";
 import { installDom } from "../../../../tests/ui/dom";
 import { ThemeProvider } from "@/browser/contexts/ThemeContext";
 import { MessageListProvider } from "./MessageListContext";
-import { getCurrentHighlightedCodeBlockLines, markdownComponents } from "./MarkdownComponents";
+import type { HighlightedCodeBlockLines } from "./MarkdownComponents";
+import { resolveCodeBlockLines, markdownComponents } from "./MarkdownComponents";
+
+// Mirrors the plain-line derivation in CodeBlock (drop the trailing empty line).
+function toPlainLines(code: string): string[] {
+  return code.split("\n").filter((line, idx, arr) => idx < arr.length - 1 || line !== "");
+}
 
 function renderCodeBlock(
   className: string,
@@ -60,29 +66,99 @@ describe("MarkdownComponents command code blocks", () => {
 
     expect(queryByRole("button", { name: "Run command" })).toBeNull();
   });
+});
+
+describe("resolveCodeBlockLines streaming highlight", () => {
+  const highlightedLine = (text: string) => `<span>${text}</span>`;
+
+  test("renders plain text before any highlight is available", () => {
+    const code = "const a = 1;\nconst b";
+    expect(resolveCodeBlockLines(null, code, toPlainLines(code), "typescript", "dark")).toEqual([
+      { content: "const a = 1;", highlighted: false },
+      { content: "const b", highlighted: false },
+    ]);
+  });
+
+  test("fully highlights when the highlight matches the current code", () => {
+    const code = "const a = 1;";
+    const highlighted: HighlightedCodeBlockLines = {
+      code,
+      shikiLanguage: "typescript",
+      theme: "dark",
+      lines: [highlightedLine("const a = 1;")],
+    };
+    expect(
+      resolveCodeBlockLines(highlighted, code, toPlainLines(code), "typescript", "dark")
+    ).toEqual([{ content: highlightedLine("const a = 1;"), highlighted: true }]);
+  });
 
-  test("ignores highlighted lines from a previous code block revision", () => {
-    const highlighted = {
-      code: "const oldValue = 1;",
+  test("keeps finalized prefix lines highlighted while the streaming tail stays plain", () => {
+    // Previous highlight ended on a newline, so its line is complete and safe to reuse.
+    const highlighted: HighlightedCodeBlockLines = {
+      code: "const a = 1;\n",
       shikiLanguage: "typescript",
-      theme: "dark" as const,
-      lines: ["<span>highlighted old value</span>"],
+      theme: "dark",
+      lines: [highlightedLine("const a = 1;")],
     };
+    const code = "const a = 1;\nconst b = 2;";
 
-    // A streaming code fence can receive a new chunk while Shiki output for the
-    // previous chunk is still cached. The renderer should fall back to current
-    // plain text until highlight output catches up to this exact code/theme tuple.
+    // The finalized first line stays colored (no flash back to plain); only the still-
+    // growing last line renders as plain text until the next highlight lands.
     expect(
-      getCurrentHighlightedCodeBlockLines(
-        highlighted,
-        "const nextValue = 2;\nconsole.log(nextValue);",
-        "typescript",
-        "dark"
-      )
-    ).toBeNull();
+      resolveCodeBlockLines(highlighted, code, toPlainLines(code), "typescript", "dark")
+    ).toEqual([
+      { content: highlightedLine("const a = 1;"), highlighted: true },
+      { content: "const b = 2;", highlighted: false },
+    ]);
+  });
+
+  test("does not reuse the last highlighted line when it was still incomplete", () => {
+    // No trailing newline => the last highlighted line ("const b = 2") was mid-stream and
+    // must not be shown as stale text once more characters arrive.
+    const highlighted: HighlightedCodeBlockLines = {
+      code: "const a = 1;\nconst b = 2",
+      shikiLanguage: "typescript",
+      theme: "dark",
+      lines: [highlightedLine("const a = 1;"), highlightedLine("const b = 2")],
+    };
+    const code = "const a = 1;\nconst b = 22;";
+
+    expect(
+      resolveCodeBlockLines(highlighted, code, toPlainLines(code), "typescript", "dark")
+    ).toEqual([
+      { content: highlightedLine("const a = 1;"), highlighted: true },
+      { content: "const b = 22;", highlighted: false },
+    ]);
+  });
+
+  test("falls back to plain text on theme change", () => {
+    const code = "const a = 1;";
+    const highlighted: HighlightedCodeBlockLines = {
+      code,
+      shikiLanguage: "typescript",
+      theme: "dark",
+      lines: [highlightedLine("const a = 1;")],
+    };
+    expect(
+      resolveCodeBlockLines(highlighted, code, toPlainLines(code), "typescript", "light")
+    ).toEqual([{ content: "const a = 1;", highlighted: false }]);
+  });
+
+  test("falls back to plain text when the highlight is not a prefix of the current code", () => {
+    const highlighted: HighlightedCodeBlockLines = {
+      code: "const oldValue = 1;\n",
+      shikiLanguage: "typescript",
+      theme: "dark",
+      lines: [highlightedLine("const oldValue = 1;")],
+    };
+    const code = "const nextValue = 2;\nconsole.log(nextValue);";
+
     expect(
-      getCurrentHighlightedCodeBlockLines(highlighted, "const oldValue = 1;", "typescript", "dark")
-    ).toEqual(["<span>highlighted old value</span>"]);
+      resolveCodeBlockLines(highlighted, code, toPlainLines(code), "typescript", "dark")
+    ).toEqual([
+      { content: "const nextValue = 2;", highlighted: false },
+      { content: "console.log(nextValue);", highlighted: false },
+    ]);
   });
 });
 

src/browser/features/Messages/MarkdownComponents.tsx

@@ -137,21 +137,64 @@ export interface HighlightedCodeBlockLines {
   lines: string[];
 }
 
-export function getCurrentHighlightedCodeBlockLines(
+export interface ResolvedCodeBlockLine {
+  /** Either Shiki HTML (when `highlighted`) or plain source text. */
+  content: string;
+  /** When true, `content` is trusted Shiki HTML; otherwise it is plain text. */
+  highlighted: boolean;
+}
+
+/**
+ * Resolve which lines to render while async Shiki output may lag behind the current code.
+ *
+ * Streaming code fences grow one chunk at a time, so the previous highlight is almost
+ * always a prefix of the current code. Every line terminated by a newline in that prefix
+ * is final (left-to-right tokenization means a completed line's colors never change based
+ * on later lines), so we keep those lines highlighted and only fall back to plain text for
+ * the still-growing tail. This shows correct partial highlighting instead of flashing the
+ * whole block between plain and colored on every streamed chunk.
+ */
+export function resolveCodeBlockLines(
   highlighted: HighlightedCodeBlockLines | null,
   code: string,
+  plainLines: string[],
   shikiLanguage: string,
   theme: "light" | "dark"
-): string[] | null {
-  if (
-    highlighted?.code === code &&
-    highlighted.shikiLanguage === shikiLanguage &&
-    highlighted.theme === theme
-  ) {
-    return highlighted.lines;
+): ResolvedCodeBlockLine[] {
+  const asPlain = (): ResolvedCodeBlockLine[] =>
+    plainLines.map((content) => ({ content, highlighted: false }));
+
+  // No highlight yet, or it targets a different language/theme: render plain.
+  if (!highlighted) {
+    return asPlain();
+  }
+  if (highlighted.shikiLanguage !== shikiLanguage || highlighted.theme !== theme) {
+    return asPlain();
   }
 
-  return null;
+  // Highlight matches the current code exactly: fully highlighted.
+  if (highlighted.code === code) {
+    return highlighted.lines.map((content) => ({ content, highlighted: true }));
+  }
+
+  // Streaming append: keep the finalized prefix highlighted, render the tail plain.
+  if (code.startsWith(highlighted.code)) {
+    // A trailing newline means every highlighted line is complete; otherwise the last
+    // highlighted line is still growing and must not be reused (it would show stale text).
+    const finalizedCount = highlighted.code.endsWith("\n")
+      ? highlighted.lines.length
+      : highlighted.lines.length - 1;
+    const safeCount = Math.max(0, Math.min(finalizedCount, plainLines.length));
+    return plainLines.map((content, idx) =>
+      idx < safeCount
+        ? { content: highlighted.lines[idx], highlighted: true }
+        : { content, highlighted: false }
+    );
+  }
+
+  // Highlight no longer corresponds to the current code (e.g. theme/code reset): render
+  // plain until the async highlighter catches up.
+  return asPlain();
 }
 
 /**
@@ -207,24 +250,17 @@ const CodeBlock: React.FC<CodeBlockProps> = ({ code, language, highlightLanguage
     ? normalizeSuggestedShellCommand(code)
     : "";
   const showRunButton = Boolean(openTerminal) && runnableCommand.length > 0;
-  // Ignore highlighted output from a previous stream chunk until the async highlighter
-  // catches up to the current code/theme. Otherwise streaming code fences briefly render
-  // stale highlighted lines with the old height, then flash to the current content.
-  const highlightedLines = getCurrentHighlightedCodeBlockLines(
-    highlighted,
-    code,
-    shikiLanguage,
-    theme
-  );
-  const lines = highlightedLines ?? plainLines;
+  // Keep finalized lines highlighted while the async highlighter catches up to the
+  // streaming tail, instead of flashing the whole block between plain and colored.
+  const lines = resolveCodeBlockLines(highlighted, code, plainLines, shikiLanguage, theme);
   const isSingleLine = lines.length === 1;
 
   return (
     <div
       className={`code-block-wrapper${isSingleLine ? " code-block-single-line" : ""}${showRunButton ? " code-block-runnable" : ""}`}
     >
       <div className="code-block-container">
-        {lines.map((content, idx) => (
+        {lines.map((line, idx) => (
           <React.Fragment key={idx}>
             <div className="line-number">{idx + 1}</div>
             {/* SECURITY AUDIT: dangerouslySetInnerHTML usage
@@ -238,9 +274,9 @@ const CodeBlock: React.FC<CodeBlockProps> = ({ code, language, highlightLanguage
              */}
             <div
               className="code-line"
-              {...(highlightedLines
-                ? { dangerouslySetInnerHTML: { __html: content } }
-                : { children: <code>{content}</code> })}
+              {...(line.highlighted
+                ? { dangerouslySetInnerHTML: { __html: line.content } }
+                : { children: <code>{line.content}</code> })}
             />
           </React.Fragment>
         ))}