From 09be92c26deeda58cff658add6f58548a75df8de Mon Sep 17 00:00:00 2001
From: stevenjj33 <75509501+stevenjj33@users.noreply.github.com>
Date: Fri, 23 Oct 2026 08:10:02 +0800
Subject: [PATCH] fix(webui): repaint mermaid diagrams on theme switch
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Review follow-up on #224 (5 comments), plus one adjacent finding.

1. A diagram already on screen did not follow a theme switch: the render
   effect read the theme from the DOM port inside the effect, whose
   dependencies were only [dom, text]. The app has no reactive theme
   store — every theme transition ends as the appearance effect flipping
   the dark class on the document element — so the port grows
   subscribeDocumentTheme (a class-attribute MutationObserver) and the
   theme joins the render dependencies. A redraw keeps the previous SVG
   on screen until the new palette is ready, so the fence never flashes
   its source text. Browser case switches theme through the real
   settings UI with a no-reload witness.

2. The theme option's doc said 'defaults to the app's current theme'
   while the code defaults to light; the comment now says what the code
   does and why (pure function, no document).

3. Documented why the retained <style> element is safe, with the probe
   evidence against the pinned library under securityLevel strict:
   classDef/style/linkStyle reject url() and stray braces at parse time,
   init directive colours go through a colour parser, and the directive
   cannot flip securityLevel/htmlLabels. Added @import / off-page url()
   stripping as the backstop for a future raw-CSS surface.

4. Adjacent finding: an @{img: 'https://...'} node makes the library
   fetch the image over the network mid-render (verified: the request
   leaves the page) — a render must not turn model output into a request
   to an arbitrary host. Such bodies now reject before the library is
   imported; data: URIs draw offline and still pass.
---
 .../features/transcript/ui/webui-mermaid.tsx  | 90 ++++++++++++++++---
 .../webui/src/client/platform/browser-dom.ts  | 25 ++++++
 .../webui/test/browser/transcript.spec.mjs    | 38 ++++++++
 .../webui/test/unit/markdown-mermaid.test.tsx | 47 ++++++++++
 4 files changed, 186 insertions(+), 14 deletions(-)

diff --git a/packages/webui/src/client/features/transcript/ui/webui-mermaid.tsx b/packages/webui/src/client/features/transcript/ui/webui-mermaid.tsx
index d5fbfa0b7..28d5f2231 100644
--- a/packages/webui/src/client/features/transcript/ui/webui-mermaid.tsx
+++ b/packages/webui/src/client/features/transcript/ui/webui-mermaid.tsx
@@ -26,8 +26,9 @@
 //      runs effects, so the block carries each state as a plain prop and the
 //      wrapper only wires the async render into it.
 //   4. The rendered SVG is sanitized before it is injected, and the diagram is
-//      drawn with the theme the app is actually showing. Both are load-bearing
-//      — see the two sections below.
+//      drawn with the theme the app is actually showing — and redrawn when
+//      that theme changes, because the palette is baked into the SVG. Both
+//      are load-bearing — see the two sections below.
 
 import { useEffect, useState, type ReactElement } from "react";
 
@@ -101,6 +102,27 @@ export function readWebuiDocumentTheme(
  * body, not assembled from user HTML, so the surface is small — but small is
  * not zero, and the filter is what makes "the SVG carries no script" a
  * property of this component rather than a property of the library version.
+ *
+ * What is deliberately kept: the `<style>` element the library emits. The
+ * diagram's class CSS lives there — every colour the theme rules assign — so
+ * removing it would unstyle the drawing. It is acceptable input because each
+ * way a diagram body can reach into it is validated by the library, checked
+ * against the pinned version under `securityLevel: "strict"`:
+ *
+ *   - `classDef` / `style` / `linkStyle` reject `url(…)` and stray `{`/`;` at
+ *     parse time (`classDef ex fill:red,stroke:url(https://…)` is a parse
+ *     error, not a rendered rule);
+ *   - the colours an `%%{init} …%%` directive can set go through a colour
+ *     parser that rejects anything that is not a colour — a breakout payload
+ *     like `#f00; } * { background-image: url(https://…)` dies as
+ *     "Unsupported color format";
+ *   - the directive cannot flip security-relevant keys at all (`securityLevel`
+ *     and `htmlLabels` stay as configured here), and `click … call` binds
+ *     nothing on the `render()` path.
+ *
+ * The `@import` and off-page `url()` rules below are the backstop: they cost
+ * nothing on a healthy diagram and they hold the line if a future diagram
+ * type grows a raw-CSS surface.
  */
 export function sanitizeWebuiMermaidSvg(svg: string): string {
   return svg
@@ -120,7 +142,13 @@ export function sanitizeWebuiMermaidSvg(svg: string): string {
     // Script-ish URLs in href / xlink:href / src.
     .replace(/\s(xlink:href|href|src)\s*=\s*"\s*javascript:[^"]*"/giu, "")
     .replace(/\s(xlink:href|href|src)\s*=\s*'\s*javascript:[^']*'/giu, "")
-    .replace(/\s(xlink:href|href|src)\s*=\s*javascript:[^\s>]*/giu, "");
+    .replace(/\s(xlink:href|href|src)\s*=\s*javascript:[^\s>]*/giu, "")
+    // CSS imports and off-page url() references. A <style> inside inline SVG
+    // applies to the whole page, not just the diagram, and an off-page url()
+    // turns a diagram into a network fetch. Fragment refs (`url(#…)` — the
+    // arrowhead markers) and data: images are drawing, not fetches, and stay.
+    .replace(/@import\b[^;{]*[;{]?/giu, "")
+    .replace(/url\(\s*(['"]?)\s*(?:https?:|\/\/|data:text\/html)[^)'"]*\1\s*\)/giu, "none");
 }
 
 /* Rendering -------------------------------------------------------------- */
@@ -147,7 +175,9 @@ async function loadMermaid(theme: WebuiMermaidTheme): Promise<MermaidModule> {
     // `theme` is the app's own theme, not the library's default: Mermaid
     // bakes its palette into the SVG it returns, so a diagram drawn with the
     // light palette is a light diagram even on a dark page. Re-initializing
-    // on a theme change is what lets the next diagram follow the switch.
+    // on a theme change is what lets the renderer follow the switch — the
+    // next diagram, and (through the subscription in `WebuiMermaid`) the
+    // ones already on screen.
     module.default.initialize({
       startOnLoad: false,
       securityLevel: "strict",
@@ -172,21 +202,43 @@ export interface WebuiMermaidRenderOptions {
    * callers never pass it.
    */
   readonly timeoutMs?: number;
-  /** The theme to draw with; defaults to the app's current theme. */
+  /**
+   * The theme to draw with. Defaults to `light`: this function is pure and
+   * has no document to ask, so an unspecified theme means the library's own
+   * default palette. The component passes the document's current theme
+   * explicitly — nothing on the production path relies on the default.
+   */
   readonly theme?: WebuiMermaidTheme;
 }
 
+/**
+ * A diagram body whose `@{img: '…'}` node image is not a `data:` URI.
+ *
+ * A remote value — an `https://` URL, the `iconify:` scheme — makes the
+ * library fetch the image over the network while it renders (verified: the
+ * request leaves the page before `render` settles). A diagram body is model
+ * output that can be steered by whatever the model read, and a render must
+ * not turn it into a request to an arbitrary host, so such a body is refused
+ * before the library is even imported. `data:` URIs are self-contained and
+ * draw offline; they pass.
+ */
+const REMOTE_IMAGE_IN_BODY = /@\{img\s*:\s*['"]\s*(?!data:)/u;
+
 /**
  * Render one Mermaid diagram body to an SVG string.
  *
  * Rejects — rather than returning a fallback — on every failure so the caller
- * decides what the user sees: an empty body, invalid syntax, a thrown render
- * and a render that outlives the timeout all end here. The empty-body check
- * runs before the dynamic import so an empty fence never pulls the library
- * down the wire.
+ * decides what the user sees: an empty body, a body that references a remote
+ * image, invalid syntax, a thrown render and a render that outlives the
+ * timeout all end here. The empty-body and remote-image checks run before the
+ * dynamic import so neither ever pulls the library down the wire, and the
+ * remote-image one also never lets it touch the network.
  */
 export async function renderWebuiMermaidSvg(text: string, options?: WebuiMermaidRenderOptions): Promise<string> {
   if (text.trim() === "") throw new Error("webui-mermaid: empty diagram body");
+  if (REMOTE_IMAGE_IN_BODY.test(text)) {
+    throw new Error("webui-mermaid: diagram body references a remote image");
+  }
   const theme = options?.theme ?? "light";
   const mermaid = (await loadMermaid(theme)).default;
   diagramSequence += 1;
@@ -269,16 +321,26 @@ export interface WebuiMermaidProps {
  * The effectful half: renders the diagram once per body and repaints the block
  * with the outcome. `cancelled` guards the state updates so a fence that
  * unmounts mid-render — a session switch, a trimmed transcript — writes
- * nothing. The dependency is the body alone: a re-render with the same text
- * must not re-run the library.
+ * nothing. The dependencies are the body and the theme: a re-render with the
+ * same text must not re-run the library, while a theme switch must — the
+ * palette is baked into the SVG, so the only way a drawn diagram follows the
+ * app's theme is to draw it again. The theme arrives as state from the DOM
+ * port's subscription (see `subscribeDocumentTheme`); until the document
+ * first changes theme after mount it is `undefined`, and the draw reads the
+ * document directly instead, so the first diagram does not wait for an event.
+ * While a theme switch redraws, the previous SVG stays on screen — a reader
+ * who changed the theme never sees the fence flash back to its source text.
  */
 export function WebuiMermaid({ text, lang }: WebuiMermaidProps): ReactElement {
   const [state, setState] = useState<WebuiMermaidState>({ status: "pending" });
   const { dom } = useWebuiBrowserCapabilities();
+  const [theme, setTheme] = useState<WebuiMermaidTheme | undefined>(undefined);
+  useEffect(() => dom.subscribeDocumentTheme(setTheme), [dom]);
   useEffect(() => {
     let cancelled = false;
-    const theme = readWebuiDocumentTheme(() => dom.readDocumentTheme());
-    renderWebuiMermaidSvg(text, { theme })
+    renderWebuiMermaidSvg(text, {
+      theme: theme ?? readWebuiDocumentTheme(() => dom.readDocumentTheme()),
+    })
       .then((svg) => {
         if (!cancelled) setState({ status: "rendered", svg });
       })
@@ -288,6 +350,6 @@ export function WebuiMermaid({ text, lang }: WebuiMermaidProps): ReactElement {
     return () => {
       cancelled = true;
     };
-  }, [dom, text]);
+  }, [dom, text, theme]);
   return <WebuiMermaidBlock state={state} text={text} lang={lang} />;
 }
diff --git a/packages/webui/src/client/platform/browser-dom.ts b/packages/webui/src/client/platform/browser-dom.ts
index 40f21f4de..369074c83 100644
--- a/packages/webui/src/client/platform/browser-dom.ts
+++ b/packages/webui/src/client/platform/browser-dom.ts
@@ -77,6 +77,20 @@ export interface WebuiBrowserDom {
    * confine to the platform layer.
    */
   readonly readDocumentTheme: () => "dark" | "light";
+  /**
+   * The theme the document is showing changed.
+   *
+   * Every theme transition the app makes ends as the appearance effect above
+   * flipping the `dark` class on the document element — a click in the
+   * settings modal, a `system` preference following the OS — so the
+   * subscription observes that class rather than any one writer: whoever
+   * flips it, the listener hears. A component whose drawing bakes the theme
+   * in (the transcript's Mermaid renderer) repaints on this signal instead of
+   * showing a stale palette until the next navigation. A platform with no
+   * document, or no `MutationObserver`, answers with a no-op: the listener
+   * keeps the theme it last read.
+   */
+  readonly subscribeDocumentTheme: (listener: (theme: "dark" | "light") => void) => () => void;
   readonly dispatchContextWindowUsageChange: () => void;
   readonly copyTextFallback: (text: string) => void;
   /**
@@ -176,6 +190,17 @@ export function createWebuiBrowserDom(
     },
     readDocumentTheme: () =>
       documentObject?.documentElement.classList.contains("dark") ? "dark" : "light",
+    subscribeDocumentTheme: (listener) => {
+      if (!documentObject || typeof MutationObserver === "undefined") return () => undefined;
+      const read = () =>
+        documentObject.documentElement.classList.contains("dark") ? "dark" : "light";
+      const observer = new MutationObserver(() => listener(read()));
+      observer.observe(documentObject.documentElement, {
+        attributes: true,
+        attributeFilter: ["class"],
+      });
+      return () => observer.disconnect();
+    },
     dispatchContextWindowUsageChange: () => {
       if (windowObject) windowObject.dispatchEvent(new Event("webui-context-window-usage-change"));
     },
diff --git a/packages/webui/test/browser/transcript.spec.mjs b/packages/webui/test/browser/transcript.spec.mjs
index 0c7e2a270..1fdc564e4 100644
--- a/packages/webui/test/browser/transcript.spec.mjs
+++ b/packages/webui/test/browser/transcript.spec.mjs
@@ -313,6 +313,44 @@ test("a diagram follows the app's theme", async ({ page }) => {
   expect(await readPalette()).not.toBe(lightPalette);
 });
 
+test("a drawn diagram repaints when the theme switches, without a reload", async ({ page }) => {
+  await stageTranscript(page, [
+    { msgId: "m-6", role: "assistant", msgContent: "```mermaid\n" + MERMAID_DIAGRAM + "\n```", timestamp: 1_700_000_000_001 },
+  ]);
+  await openApp(page, "#session=A");
+  await expect(page.locator(".webui-mermaid svg")).toBeVisible({ timeout: 20_000 });
+  const readPalette = () =>
+    page.locator(".webui-mermaid svg").evaluate((node) => node.querySelector("style")?.textContent ?? "");
+  const lightPalette = await readPalette();
+
+  // No reload happens anywhere below this marker — a reload would wipe it, so
+  // reading it back at the end is the proof that the repaint came from the
+  // renderer hearing the theme change, not from the page coming back.
+  await page.evaluate(() => {
+    window.__webuiNoReloadWitness = true;
+  });
+
+  // The app's own control, driven the way a user drives it (the same flow
+  // `settings-theme.spec.mjs` uses): user menu, 设置, 深色模式. That click
+  // flips the `dark` class on `<html>`, and the diagram must redraw — the
+  // palette is baked into the SVG, so a theme switch that does not redraw
+  // leaves a light diagram sitting on a dark page until the next navigation.
+  // This is the path the reload-based case above cannot see.
+  await page.locator('[data-testid="sidebar-user-menu-trigger"]').click();
+  await page.locator('[data-testid="user-menu-settings"]').click();
+  await expect(page.locator('[role="dialog"][aria-label="设置"]')).toBeVisible();
+  await page.locator('[data-testid="mavis-settings-appearance-row"] .mavis-settings-theme-option', { hasText: "深色模式" }).click();
+  await expect(page.locator("html")).toHaveClass(/dark/);
+
+  await expect
+    .poll(async () => readPalette(), { timeout: 20_000 })
+    .not.toBe(lightPalette);
+  expect(await page.evaluate(() => window.__webuiNoReloadWitness)).toBe(true);
+  // The repaint produced a rendered diagram again — not a fallback.
+  await expect(page.locator(".webui-mermaid svg")).toBeVisible();
+  await expect(page.locator(".webui-code-block")).toHaveCount(0);
+});
+
 test("a wide diagram fits its container", async ({ page }) => {
   // A wide diagram must not push the transcript sideways: the SVG scales with
   // the container (`max-width: 100%`, `height: auto`).
diff --git a/packages/webui/test/unit/markdown-mermaid.test.tsx b/packages/webui/test/unit/markdown-mermaid.test.tsx
index d69fed57a..0c103618b 100644
--- a/packages/webui/test/unit/markdown-mermaid.test.tsx
+++ b/packages/webui/test/unit/markdown-mermaid.test.tsx
@@ -129,6 +129,31 @@ describe("renderWebuiMermaidSvg", () => {
     await expect(renderWebuiMermaidSvg("this is not a diagram @@")).rejects.toThrow(/Parse error/iu);
   });
 
+  it("rejects a body that references a remote image, without touching the library", async () => {
+    // `@{img: 'https://…'}` makes the real library fetch the image over the
+    // network mid-render (the request leaves the page before `render`
+    // settles), and `iconify:` does the same through its icon loader. A
+    // diagram body is model output, so the render must end before it starts
+    // — this path is the plain-fence fallback, same as any other refusal.
+    for (const body of [
+      "graph TD;\n  A@{img: 'https://evil.example/logo.png'} --> B;",
+      'graph TD;\n  A@{img: "iconify:mdi:account"} --> B;',
+    ]) {
+      await expect(renderWebuiMermaidSvg(body)).rejects.toThrow(/remote image/iu);
+    }
+    expect(renderMock).not.toHaveBeenCalled();
+  });
+
+  it("draws a self-contained data-URI image", async () => {
+    // `data:` URIs carry their bytes with them and draw offline — the guard
+    // above must not take the one image form that never touches the network.
+    renderMock.mockResolvedValue({ svg: '<svg id="webui-mermaid-2"><g></g></svg>' } as never);
+    await expect(
+      renderWebuiMermaidSvg("graph TD;\n  A@{img: 'data:image/svg+xml;base64,PHN2Zw=='} --> B;"),
+    ).resolves.toContain("<svg");
+    expect(renderMock).toHaveBeenCalled();
+  });
+
   it("rejects a render that outlives the timeout", async () => {
     // A real clock with an injected budget: faking timers here would put the
     // rejection inside a timer-advance window, which Vitest reports as an
@@ -272,6 +297,28 @@ describe("the SVG that reaches the DOM", () => {
     expect(clean).toContain("<g/>");
   });
 
+  it("neutralises CSS imports and off-page url() references, keeping fragment refs and data images", () => {
+    // The `<style>` element is kept (it carries the diagram's class CSS);
+    // these rules are its backstop. Fragment refs (`url(#…)` — the arrowhead
+    // markers) and data: images are drawing, not fetches, and must survive.
+    const styled = [
+      "<svg><style>",
+      ".edge { fill: red; }",
+      '@import url("https://evil.example/i.css");',
+      ".x { background: url(https://evil.example/b.png); }",
+      '.y { background: url("//evil.example/p.png"); }',
+      ".marker { fill: url(#arrow); }",
+      '.z { background: url("data:image/svg+xml;base64,PHN2Zw=="); }',
+      "</style></svg>",
+    ].join("");
+    const clean = sanitizeWebuiMermaidSvg(styled);
+    expect(clean).not.toContain("@import");
+    expect(clean).not.toContain("evil.example");
+    expect(clean).toContain(".edge { fill: red; }");
+    expect(clean).toContain("url(#arrow)");
+    expect(clean).toContain("data:image/svg+xml;base64,PHN2Zw==");
+  });
+
   it("leaves an ordinary diagram untouched", () => {
     const plain = '<svg class="flowchart" viewBox="0 0 10 10"><g class="node"><text>开始</text></g></svg>';
     expect(sanitizeWebuiMermaidSvg(plain)).toBe(plain);