Skip to content

Commit 2bd11d0

Browse files
authored
Merge pull request #231 from mcode-webui/fix/webui-mermaid-theme-rerender
fix(webui): Mermaid 已画图形随主题切换重画(#224 评审跟进)
2 parents e468be0 + 09be92c commit 2bd11d0

4 files changed

Lines changed: 186 additions & 14 deletions

File tree

‎packages/webui/src/client/features/transcript/ui/webui-mermaid.tsx‎

Lines changed: 76 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -26,8 +26,9 @@
2626
// runs effects, so the block carries each state as a plain prop and the
2727
// wrapper only wires the async render into it.
2828
// 4. The rendered SVG is sanitized before it is injected, and the diagram is
29-
// drawn with the theme the app is actually showing. Both are load-bearing
30-
// — see the two sections below.
29+
// drawn with the theme the app is actually showing — and redrawn when
30+
// that theme changes, because the palette is baked into the SVG. Both
31+
// are load-bearing — see the two sections below.
3132

3233
import { useEffect, useState, type ReactElement } from "react";
3334

@@ -101,6 +102,27 @@ export function readWebuiDocumentTheme(
101102
* body, not assembled from user HTML, so the surface is small — but small is
102103
* not zero, and the filter is what makes "the SVG carries no script" a
103104
* property of this component rather than a property of the library version.
105+
*
106+
* What is deliberately kept: the `<style>` element the library emits. The
107+
* diagram's class CSS lives there — every colour the theme rules assign — so
108+
* removing it would unstyle the drawing. It is acceptable input because each
109+
* way a diagram body can reach into it is validated by the library, checked
110+
* against the pinned version under `securityLevel: "strict"`:
111+
*
112+
* - `classDef` / `style` / `linkStyle` reject `url(…)` and stray `{`/`;` at
113+
* parse time (`classDef ex fill:red,stroke:url(https://…)` is a parse
114+
* error, not a rendered rule);
115+
* - the colours an `%%{init} …%%` directive can set go through a colour
116+
* parser that rejects anything that is not a colour — a breakout payload
117+
* like `#f00; } * { background-image: url(https://…)` dies as
118+
* "Unsupported color format";
119+
* - the directive cannot flip security-relevant keys at all (`securityLevel`
120+
* and `htmlLabels` stay as configured here), and `click … call` binds
121+
* nothing on the `render()` path.
122+
*
123+
* The `@import` and off-page `url()` rules below are the backstop: they cost
124+
* nothing on a healthy diagram and they hold the line if a future diagram
125+
* type grows a raw-CSS surface.
104126
*/
105127
export function sanitizeWebuiMermaidSvg(svg: string): string {
106128
return svg
@@ -120,7 +142,13 @@ export function sanitizeWebuiMermaidSvg(svg: string): string {
120142
// Script-ish URLs in href / xlink:href / src.
121143
.replace(/\s(xlink:href|href|src)\s*=\s*"\s*javascript:[^"]*"/giu, "")
122144
.replace(/\s(xlink:href|href|src)\s*=\s*'\s*javascript:[^']*'/giu, "")
123-
.replace(/\s(xlink:href|href|src)\s*=\s*javascript:[^\s>]*/giu, "");
145+
.replace(/\s(xlink:href|href|src)\s*=\s*javascript:[^\s>]*/giu, "")
146+
// CSS imports and off-page url() references. A <style> inside inline SVG
147+
// applies to the whole page, not just the diagram, and an off-page url()
148+
// turns a diagram into a network fetch. Fragment refs (`url(#…)` — the
149+
// arrowhead markers) and data: images are drawing, not fetches, and stay.
150+
.replace(/@import\b[^;{]*[;{]?/giu, "")
151+
.replace(/url\(\s*(['"]?)\s*(?:https?:|\/\/|data:text\/html)[^)'"]*\1\s*\)/giu, "none");
124152
}
125153

126154
/* Rendering -------------------------------------------------------------- */
@@ -147,7 +175,9 @@ async function loadMermaid(theme: WebuiMermaidTheme): Promise<MermaidModule> {
147175
// `theme` is the app's own theme, not the library's default: Mermaid
148176
// bakes its palette into the SVG it returns, so a diagram drawn with the
149177
// light palette is a light diagram even on a dark page. Re-initializing
150-
// on a theme change is what lets the next diagram follow the switch.
178+
// on a theme change is what lets the renderer follow the switch — the
179+
// next diagram, and (through the subscription in `WebuiMermaid`) the
180+
// ones already on screen.
151181
module.default.initialize({
152182
startOnLoad: false,
153183
securityLevel: "strict",
@@ -172,21 +202,43 @@ export interface WebuiMermaidRenderOptions {
172202
* callers never pass it.
173203
*/
174204
readonly timeoutMs?: number;
175-
/** The theme to draw with; defaults to the app's current theme. */
205+
/**
206+
* The theme to draw with. Defaults to `light`: this function is pure and
207+
* has no document to ask, so an unspecified theme means the library's own
208+
* default palette. The component passes the document's current theme
209+
* explicitly — nothing on the production path relies on the default.
210+
*/
176211
readonly theme?: WebuiMermaidTheme;
177212
}
178213

214+
/**
215+
* A diagram body whose `@{img: '…'}` node image is not a `data:` URI.
216+
*
217+
* A remote value — an `https://` URL, the `iconify:` scheme — makes the
218+
* library fetch the image over the network while it renders (verified: the
219+
* request leaves the page before `render` settles). A diagram body is model
220+
* output that can be steered by whatever the model read, and a render must
221+
* not turn it into a request to an arbitrary host, so such a body is refused
222+
* before the library is even imported. `data:` URIs are self-contained and
223+
* draw offline; they pass.
224+
*/
225+
const REMOTE_IMAGE_IN_BODY = /@\{img\s*:\s*['"]\s*(?!data:)/u;
226+
179227
/**
180228
* Render one Mermaid diagram body to an SVG string.
181229
*
182230
* Rejects — rather than returning a fallback — on every failure so the caller
183-
* decides what the user sees: an empty body, invalid syntax, a thrown render
184-
* and a render that outlives the timeout all end here. The empty-body check
185-
* runs before the dynamic import so an empty fence never pulls the library
186-
* down the wire.
231+
* decides what the user sees: an empty body, a body that references a remote
232+
* image, invalid syntax, a thrown render and a render that outlives the
233+
* timeout all end here. The empty-body and remote-image checks run before the
234+
* dynamic import so neither ever pulls the library down the wire, and the
235+
* remote-image one also never lets it touch the network.
187236
*/
188237
export async function renderWebuiMermaidSvg(text: string, options?: WebuiMermaidRenderOptions): Promise<string> {
189238
if (text.trim() === "") throw new Error("webui-mermaid: empty diagram body");
239+
if (REMOTE_IMAGE_IN_BODY.test(text)) {
240+
throw new Error("webui-mermaid: diagram body references a remote image");
241+
}
190242
const theme = options?.theme ?? "light";
191243
const mermaid = (await loadMermaid(theme)).default;
192244
diagramSequence += 1;
@@ -269,16 +321,26 @@ export interface WebuiMermaidProps {
269321
* The effectful half: renders the diagram once per body and repaints the block
270322
* with the outcome. `cancelled` guards the state updates so a fence that
271323
* unmounts mid-render — a session switch, a trimmed transcript — writes
272-
* nothing. The dependency is the body alone: a re-render with the same text
273-
* must not re-run the library.
324+
* nothing. The dependencies are the body and the theme: a re-render with the
325+
* same text must not re-run the library, while a theme switch must — the
326+
* palette is baked into the SVG, so the only way a drawn diagram follows the
327+
* app's theme is to draw it again. The theme arrives as state from the DOM
328+
* port's subscription (see `subscribeDocumentTheme`); until the document
329+
* first changes theme after mount it is `undefined`, and the draw reads the
330+
* document directly instead, so the first diagram does not wait for an event.
331+
* While a theme switch redraws, the previous SVG stays on screen — a reader
332+
* who changed the theme never sees the fence flash back to its source text.
274333
*/
275334
export function WebuiMermaid({ text, lang }: WebuiMermaidProps): ReactElement {
276335
const [state, setState] = useState<WebuiMermaidState>({ status: "pending" });
277336
const { dom } = useWebuiBrowserCapabilities();
337+
const [theme, setTheme] = useState<WebuiMermaidTheme | undefined>(undefined);
338+
useEffect(() => dom.subscribeDocumentTheme(setTheme), [dom]);
278339
useEffect(() => {
279340
let cancelled = false;
280-
const theme = readWebuiDocumentTheme(() => dom.readDocumentTheme());
281-
renderWebuiMermaidSvg(text, { theme })
341+
renderWebuiMermaidSvg(text, {
342+
theme: theme ?? readWebuiDocumentTheme(() => dom.readDocumentTheme()),
343+
})
282344
.then((svg) => {
283345
if (!cancelled) setState({ status: "rendered", svg });
284346
})
@@ -288,6 +350,6 @@ export function WebuiMermaid({ text, lang }: WebuiMermaidProps): ReactElement {
288350
return () => {
289351
cancelled = true;
290352
};
291-
}, [dom, text]);
353+
}, [dom, text, theme]);
292354
return <WebuiMermaidBlock state={state} text={text} lang={lang} />;
293355
}

‎packages/webui/src/client/platform/browser-dom.ts‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,20 @@ export interface WebuiBrowserDom {
7777
* confine to the platform layer.
7878
*/
7979
readonly readDocumentTheme: () => "dark" | "light";
80+
/**
81+
* The theme the document is showing changed.
82+
*
83+
* Every theme transition the app makes ends as the appearance effect above
84+
* flipping the `dark` class on the document element — a click in the
85+
* settings modal, a `system` preference following the OS — so the
86+
* subscription observes that class rather than any one writer: whoever
87+
* flips it, the listener hears. A component whose drawing bakes the theme
88+
* in (the transcript's Mermaid renderer) repaints on this signal instead of
89+
* showing a stale palette until the next navigation. A platform with no
90+
* document, or no `MutationObserver`, answers with a no-op: the listener
91+
* keeps the theme it last read.
92+
*/
93+
readonly subscribeDocumentTheme: (listener: (theme: "dark" | "light") => void) => () => void;
8094
readonly dispatchContextWindowUsageChange: () => void;
8195
readonly copyTextFallback: (text: string) => void;
8296
/**
@@ -176,6 +190,17 @@ export function createWebuiBrowserDom(
176190
},
177191
readDocumentTheme: () =>
178192
documentObject?.documentElement.classList.contains("dark") ? "dark" : "light",
193+
subscribeDocumentTheme: (listener) => {
194+
if (!documentObject || typeof MutationObserver === "undefined") return () => undefined;
195+
const read = () =>
196+
documentObject.documentElement.classList.contains("dark") ? "dark" : "light";
197+
const observer = new MutationObserver(() => listener(read()));
198+
observer.observe(documentObject.documentElement, {
199+
attributes: true,
200+
attributeFilter: ["class"],
201+
});
202+
return () => observer.disconnect();
203+
},
179204
dispatchContextWindowUsageChange: () => {
180205
if (windowObject) windowObject.dispatchEvent(new Event("webui-context-window-usage-change"));
181206
},

‎packages/webui/test/browser/transcript.spec.mjs‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -413,6 +413,44 @@ test("a diagram follows the app's theme", async ({ page }) => {
413413
expect(await readPalette()).not.toBe(lightPalette);
414414
});
415415

416+
test("a drawn diagram repaints when the theme switches, without a reload", async ({ page }) => {
417+
await stageTranscript(page, [
418+
{ msgId: "m-6", role: "assistant", msgContent: "```mermaid\n" + MERMAID_DIAGRAM + "\n```", timestamp: 1_700_000_000_001 },
419+
]);
420+
await openApp(page, "#session=A");
421+
await expect(page.locator(".webui-mermaid svg")).toBeVisible({ timeout: 20_000 });
422+
const readPalette = () =>
423+
page.locator(".webui-mermaid svg").evaluate((node) => node.querySelector("style")?.textContent ?? "");
424+
const lightPalette = await readPalette();
425+
426+
// No reload happens anywhere below this marker — a reload would wipe it, so
427+
// reading it back at the end is the proof that the repaint came from the
428+
// renderer hearing the theme change, not from the page coming back.
429+
await page.evaluate(() => {
430+
window.__webuiNoReloadWitness = true;
431+
});
432+
433+
// The app's own control, driven the way a user drives it (the same flow
434+
// `settings-theme.spec.mjs` uses): user menu, 设置, 深色模式. That click
435+
// flips the `dark` class on `<html>`, and the diagram must redraw — the
436+
// palette is baked into the SVG, so a theme switch that does not redraw
437+
// leaves a light diagram sitting on a dark page until the next navigation.
438+
// This is the path the reload-based case above cannot see.
439+
await page.locator('[data-testid="sidebar-user-menu-trigger"]').click();
440+
await page.locator('[data-testid="user-menu-settings"]').click();
441+
await expect(page.locator('[role="dialog"][aria-label="设置"]')).toBeVisible();
442+
await page.locator('[data-testid="mavis-settings-appearance-row"] .mavis-settings-theme-option', { hasText: "深色模式" }).click();
443+
await expect(page.locator("html")).toHaveClass(/dark/);
444+
445+
await expect
446+
.poll(async () => readPalette(), { timeout: 20_000 })
447+
.not.toBe(lightPalette);
448+
expect(await page.evaluate(() => window.__webuiNoReloadWitness)).toBe(true);
449+
// The repaint produced a rendered diagram again — not a fallback.
450+
await expect(page.locator(".webui-mermaid svg")).toBeVisible();
451+
await expect(page.locator(".webui-code-block")).toHaveCount(0);
452+
});
453+
416454
test("a wide diagram fits its container", async ({ page }) => {
417455
// A wide diagram must not push the transcript sideways: the SVG scales with
418456
// the container (`max-width: 100%`, `height: auto`).

‎packages/webui/test/unit/markdown-mermaid.test.tsx‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -129,6 +129,31 @@ describe("renderWebuiMermaidSvg", () => {
129129
await expect(renderWebuiMermaidSvg("this is not a diagram @@")).rejects.toThrow(/Parse error/iu);
130130
});
131131

132+
it("rejects a body that references a remote image, without touching the library", async () => {
133+
// `@{img: 'https://…'}` makes the real library fetch the image over the
134+
// network mid-render (the request leaves the page before `render`
135+
// settles), and `iconify:` does the same through its icon loader. A
136+
// diagram body is model output, so the render must end before it starts
137+
// — this path is the plain-fence fallback, same as any other refusal.
138+
for (const body of [
139+
"graph TD;\n A@{img: 'https://evil.example/logo.png'} --> B;",
140+
'graph TD;\n A@{img: "iconify:mdi:account"} --> B;',
141+
]) {
142+
await expect(renderWebuiMermaidSvg(body)).rejects.toThrow(/remote image/iu);
143+
}
144+
expect(renderMock).not.toHaveBeenCalled();
145+
});
146+
147+
it("draws a self-contained data-URI image", async () => {
148+
// `data:` URIs carry their bytes with them and draw offline — the guard
149+
// above must not take the one image form that never touches the network.
150+
renderMock.mockResolvedValue({ svg: '<svg id="webui-mermaid-2"><g></g></svg>' } as never);
151+
await expect(
152+
renderWebuiMermaidSvg("graph TD;\n A@{img: 'data:image/svg+xml;base64,PHN2Zw=='} --> B;"),
153+
).resolves.toContain("<svg");
154+
expect(renderMock).toHaveBeenCalled();
155+
});
156+
132157
it("rejects a render that outlives the timeout", async () => {
133158
// A real clock with an injected budget: faking timers here would put the
134159
// rejection inside a timer-advance window, which Vitest reports as an
@@ -272,6 +297,28 @@ describe("the SVG that reaches the DOM", () => {
272297
expect(clean).toContain("<g/>");
273298
});
274299

300+
it("neutralises CSS imports and off-page url() references, keeping fragment refs and data images", () => {
301+
// The `<style>` element is kept (it carries the diagram's class CSS);
302+
// these rules are its backstop. Fragment refs (`url(#…)` — the arrowhead
303+
// markers) and data: images are drawing, not fetches, and must survive.
304+
const styled = [
305+
"<svg><style>",
306+
".edge { fill: red; }",
307+
'@import url("https://evil.example/i.css");',
308+
".x { background: url(https://evil.example/b.png); }",
309+
'.y { background: url("//evil.example/p.png"); }',
310+
".marker { fill: url(#arrow); }",
311+
'.z { background: url("data:image/svg+xml;base64,PHN2Zw=="); }',
312+
"</style></svg>",
313+
].join("");
314+
const clean = sanitizeWebuiMermaidSvg(styled);
315+
expect(clean).not.toContain("@import");
316+
expect(clean).not.toContain("evil.example");
317+
expect(clean).toContain(".edge { fill: red; }");
318+
expect(clean).toContain("url(#arrow)");
319+
expect(clean).toContain("data:image/svg+xml;base64,PHN2Zw==");
320+
});
321+
275322
it("leaves an ordinary diagram untouched", () => {
276323
const plain = '<svg class="flowchart" viewBox="0 0 10 10"><g class="node"><text>开始</text></g></svg>';
277324
expect(sanitizeWebuiMermaidSvg(plain)).toBe(plain);

0 commit comments

Comments
 (0)