From ed91aa31ff8ea3810b64f4abf6b4887815ed0c46 Mon Sep 17 00:00:00 2001 From: Ammar Date: Sat, 6 Jun 2026 10:55:22 -0500 Subject: [PATCH] fix: keep finalized code lines highlighted while streaming --- .../Messages/MarkdownComponents.test.tsx | 112 +++++++++++++++--- .../features/Messages/MarkdownComponents.tsx | 82 +++++++++---- 2 files changed, 153 insertions(+), 41 deletions(-) diff --git a/src/browser/features/Messages/MarkdownComponents.test.tsx b/src/browser/features/Messages/MarkdownComponents.test.tsx index 3636d04d14..a17b916bed 100644 --- a/src/browser/features/Messages/MarkdownComponents.test.tsx +++ b/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) => `${text}`; + + 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: ["highlighted old value"], + 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(["highlighted old value"]); + resolveCodeBlockLines(highlighted, code, toPlainLines(code), "typescript", "dark") + ).toEqual([ + { content: "const nextValue = 2;", highlighted: false }, + { content: "console.log(nextValue);", highlighted: false }, + ]); }); }); diff --git a/src/browser/features/Messages/MarkdownComponents.tsx b/src/browser/features/Messages/MarkdownComponents.tsx index 6196ce2ae7..e65a2c678e 100644 --- a/src/browser/features/Messages/MarkdownComponents.tsx +++ b/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,16 +250,9 @@ const CodeBlock: React.FC = ({ 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 ( @@ -224,7 +260,7 @@ const CodeBlock: React.FC = ({ code, language, highlightLanguage className={`code-block-wrapper${isSingleLine ? " code-block-single-line" : ""}${showRunButton ? " code-block-runnable" : ""}`} >
- {lines.map((content, idx) => ( + {lines.map((line, idx) => (
{idx + 1}
{/* SECURITY AUDIT: dangerouslySetInnerHTML usage @@ -238,9 +274,9 @@ const CodeBlock: React.FC = ({ code, language, highlightLanguage */}
{content} })} + {...(line.highlighted + ? { dangerouslySetInnerHTML: { __html: line.content } } + : { children: {line.content} })} /> ))}