fix(ui): preserve user code block rendering (#85942)

ui/src/ui/chat/grouped-render.test.ts

@@ -12,6 +12,9 @@ import {
 import { normalizeMessage } from "./message-normalizer.ts";
 
 const localStorageValues = vi.hoisted(() => new Map<string, string>());
+const markdownRenderMock = vi.hoisted(() =>
+  vi.fn((value: string, _options?: { codeBlockChrome?: "copy" | "none" }) => value),
+);
 
 vi.mock("../../local-storage.ts", () => ({
   getSafeLocalStorage: () => ({
@@ -22,7 +25,7 @@ vi.mock("../../local-storage.ts", () => ({
 }));
 
 vi.mock("../markdown.ts", () => ({
-  toSanitizedMarkdownHtml: (value: string) => value,
+  toSanitizedMarkdownHtml: markdownRenderMock,
 }));
 
 vi.mock("../icons.ts", () => ({
@@ -494,6 +497,7 @@ function mediaTicketPayload(mediaTicket: string, ttlMs = 5 * 60 * 1000) {
 }
 
 afterEach(() => {
+  markdownRenderMock.mockClear();
   document.querySelectorAll("[data-delete-confirm-fixture]").forEach((element) => {
     element.remove();
   });
@@ -566,6 +570,36 @@ describe("grouped chat rendering", () => {
     expect(userBubble.querySelector(".chat-bubble-actions")).toBeNull();
   });
 
+  it("renders user markdown without code-block copy chrome", () => {
+    const container = document.createElement("div");
+    const markdown = "```bash\npython3 - <<'PY'\nprint('ok')\nPY\n```";
+
+    renderGroupedMessage(
+      container,
+      {
+        role: "user",
+        content: markdown,
+        timestamp: 1001,
+      },
+      "user",
+    );
+
+    expect(markdownRenderMock).toHaveBeenCalledWith(markdown, { codeBlockChrome: "none" });
+  });
+
+  it("keeps assistant markdown code-block copy chrome enabled", () => {
+    const container = document.createElement("div");
+    const markdown = "```bash\necho ok\n```";
+
+    renderAssistantMessage(container, {
+      role: "assistant",
+      content: markdown,
+      timestamp: 1000,
+    });
+
+    expect(markdownRenderMock).toHaveBeenCalledWith(markdown, undefined);
+  });
+
   it("positions delete confirm by message side", () => {
     const container = document.createElement("div");
     clearDeleteConfirmSkip();

ui/src/ui/chat/grouped-render.ts

@@ -1498,6 +1498,7 @@ function renderGroupedMessage(
   const markdownBase = extractedText?.trim() ? extractedText : null;
   const reasoningMarkdown = extractedThinking ? formatReasoningMarkdown(extractedThinking) : null;
   const markdown = markdownBase;
+  const markdownRenderOptions = role === "user" ? { codeBlockChrome: "none" as const } : undefined;
   const canCopyMarkdown = role === "assistant" && Boolean(markdown?.trim());
   const canExpand = role === "assistant" && Boolean(onOpenSidebar && markdown?.trim());
   const hasActions = canCopyMarkdown || canExpand;
@@ -1627,7 +1628,9 @@ function renderGroupedMessage(
                           </details>`
                         : markdown
                           ? html`<div class="chat-text" dir="${detectTextDirection(markdown)}">
-                              ${unsafeHTML(toSanitizedMarkdownHtml(markdown))}
+                              ${unsafeHTML(
+                                toSanitizedMarkdownHtml(markdown, markdownRenderOptions),
+                              )}
                             </div>`
                           : nothing}
                       ${hasToolCards
@@ -1689,7 +1692,7 @@ function renderGroupedMessage(
                 </details>`
               : markdown
                 ? html`<div class="chat-text" dir="${detectTextDirection(markdown)}">
-                    ${unsafeHTML(toSanitizedMarkdownHtml(markdown))}
+                    ${unsafeHTML(toSanitizedMarkdownHtml(markdown, markdownRenderOptions))}
                   </div>`
                 : nothing}
             ${hasToolCards

ui/src/ui/markdown.test.ts

@@ -367,6 +367,38 @@ describe("toSanitizedMarkdownHtml", () => {
       );
     });
 
+    it("omits copy chrome when rendering user-preserved code blocks", () => {
+      const source = `python3 - <<'PY'
+import openpyxl
+
+for ws in wb.worksheets:
+    print(f"--- {ws.title} ---")
+    rows = 0
+
+    for row in ws.iter_rows(values_only=True):
+        print(row)
+PY
+`;
+      const html = toSanitizedMarkdownHtml(`\`\`\`bash\n${source}\`\`\``, {
+        codeBlockChrome: "none",
+      });
+      const fragment = htmlFragment(html);
+
+      expect(fragment.querySelector(".code-block-copy")).toBeNull();
+      expect(fragment.textContent).toBe(source);
+    });
+
+    it("keeps the no-chrome code-block cache separate from copy-enabled rendering", () => {
+      const markdown = "```\ncode\n```";
+      const plain = toSanitizedMarkdownHtml(markdown, { codeBlockChrome: "none" });
+      const copyable = toSanitizedMarkdownHtml(markdown);
+
+      expect(htmlFragment(plain).querySelector(".code-block-copy")).toBeNull();
+      expect(htmlFragment(copyable).querySelector(".code-block-copy")).toBeInstanceOf(
+        HTMLButtonElement,
+      );
+    });
+
     it("highlights fenced code blocks while preserving copy text", () => {
       const source = 'const answer = "yes";\nconsole.log(answer);\n';
       const html = toSanitizedMarkdownHtml(`\`\`\`js\n${source}\`\`\``);

ui/src/ui/markdown.ts

@@ -87,6 +87,16 @@ const INLINE_DATA_IMAGE_RE = /^data:image\/[a-z0-9.+-]+;base64,/i;
 const markdownCache = new Map<string, string>();
 const TAIL_LINK_BLUR_CLASS = "chat-link-tail-blur";
 
+export type MarkdownCodeBlockChrome = "copy" | "none";
+
+export type MarkdownRenderOptions = {
+  codeBlockChrome?: MarkdownCodeBlockChrome;
+};
+
+type MarkdownRenderEnv = {
+  codeBlockChrome: MarkdownCodeBlockChrome;
+};
+
 // CJK character ranges for URL boundary detection (RFC 3986: CJK is not valid in raw URLs).
 // CJK Unified Ideographs, CJK Symbols/Punctuation, Fullwidth Forms, Hiragana, Katakana,
 // Hangul Syllables, and CJK Compatibility Ideographs.
@@ -115,6 +125,16 @@ function setCachedMarkdown(key: string, value: string) {
   }
 }
 
+function normalizeMarkdownRenderOptions(options: MarkdownRenderOptions = {}): MarkdownRenderEnv {
+  return {
+    codeBlockChrome: options.codeBlockChrome ?? "copy",
+  };
+}
+
+function shouldRenderCodeBlockCopy(env: unknown): boolean {
+  return (env as Partial<MarkdownRenderEnv> | undefined)?.codeBlockChrome !== "none";
+}
+
 function installHooks() {
   if (hooksInstalled) {
     return;
@@ -521,7 +541,7 @@ md.renderer.rules.image = (tokens, idx) => {
 };
 
 // Override fenced code blocks with copy button + JSON collapse
-md.renderer.rules.fence = (tokens, idx) => {
+md.renderer.rules.fence = (tokens, idx, _options, env) => {
   const token = tokens[idx];
   // token.info contains the full fence info string (e.g., "json title=foo");
   // extract only the first whitespace-separated token as the language.
@@ -530,6 +550,9 @@ md.renderer.rules.fence = (tokens, idx) => {
   const highlighted = highlightCode(text, lang);
   const classAttr = codeClassAttribute(lang, highlighted);
   const codeBlock = `<pre><code${classAttr}>${highlighted}</code></pre>`;
+  if (!shouldRenderCodeBlockCopy(env)) {
+    return codeBlock;
+  }
   const langLabel = lang ? `<span class="code-block-lang">${escapeHtml(lang)}</span>` : "";
   const attrSafe = escapeHtml(text);
   const copyBtn = `<button type="button" class="code-block-copy" data-code="${attrSafe}" aria-label="${escapeHtml(t("common.copyCode"))}"><span class="code-block-copy__idle">${escapeHtml(t("common.copy"))}</span><span class="code-block-copy__done">${escapeHtml(t("common.copied"))}</span></button>`;
@@ -552,12 +575,15 @@ md.renderer.rules.fence = (tokens, idx) => {
 };
 
 // Override indented code blocks (code_block) with the same treatment as fence
-md.renderer.rules.code_block = (tokens, idx) => {
+md.renderer.rules.code_block = (tokens, idx, _options, env) => {
   const token = tokens[idx];
   const text = token.content;
   const highlighted = highlightCode(text, "");
   const classAttr = codeClassAttribute("", highlighted);
   const codeBlock = `<pre><code${classAttr}>${highlighted}</code></pre>`;
+  if (!shouldRenderCodeBlockCopy(env)) {
+    return codeBlock;
+  }
   const attrSafe = escapeHtml(text);
   const copyBtn = `<button type="button" class="code-block-copy" data-code="${attrSafe}" aria-label="${escapeHtml(t("common.copyCode"))}"><span class="code-block-copy__idle">${escapeHtml(t("common.copy"))}</span><span class="code-block-copy__done">${escapeHtml(t("common.copied"))}</span></button>`;
   const header = `<div class="code-block-header">${copyBtn}</div>`;
@@ -576,13 +602,17 @@ md.renderer.rules.code_block = (tokens, idx) => {
   return `<div class="code-block-wrapper">${header}${codeBlock}</div>`;
 };
 
-export function toSanitizedMarkdownHtml(markdown: string): string {
+export function toSanitizedMarkdownHtml(
+  markdown: string,
+  options: MarkdownRenderOptions = {},
+): string {
+  const renderOptions = normalizeMarkdownRenderOptions(options);
   const input = stripUnsupportedCitationControlMarkers(markdown).trim();
   if (!input) {
     return "";
   }
   installHooks();
-  const cacheKey = `${i18n.getLocale()}\0${input}`;
+  const cacheKey = `${i18n.getLocale()}\0${renderOptions.codeBlockChrome}\0${input}`;
   if (input.length <= MARKDOWN_CACHE_MAX_CHARS) {
     const cached = getCachedMarkdown(cacheKey);
     if (cached !== null) {
@@ -606,7 +636,7 @@ export function toSanitizedMarkdownHtml(markdown: string): string {
   }
   let rendered: string;
   try {
-    rendered = md.render(`${truncated.text}${suffix}`);
+    rendered = md.render(`${truncated.text}${suffix}`, renderOptions);
   } catch (err) {
     // Fall back to escaped plain text when md.render() throws (#36213).
     console.warn("[markdown] md.render failed, falling back to plain text:", err);