fix(agents): render local and web images in replies (#1711)

e2e/agentMarkdown.spec.ts

@@ -47,11 +47,16 @@ const REPLY = [
   "```",
 ].join("\n");
 
-async function openReply(page: Page, text: string) {
+async function openReply(page: Page, text: string, imageData?: string) {
   await page.goto("/");
   await page.getByRole("tree", { name: "Project" }).waitFor();
   await page.evaluate(
-    async ({ task, text }) => {
+    async ({ task, text, imageData }) => {
+      if (imageData) {
+        const routePath = "/src/platform/tauriCore.ts";
+        const { setInvokeRoute } = await import(routePath);
+        setInvokeRoute((cmd: string) => (cmd === "fs_read_image" ? Promise.resolve(imageData) : null));
+      }
       // Paths go through variables: knip can't resolve these browser-absolute URLs.
       const storePath = "/src/features/agents/store.ts";
       const store = await import(storePath);
@@ -66,7 +71,7 @@ async function openReply(page: Page, text: string) {
       });
       store.dockTask(task.id);
     },
-    { task: TASK, text },
+    { task: TASK, text, imageData },
   );
   const reply = page.locator("[class*='_markdown_']").first();
   await reply.locator("h1, p").first().waitFor();
@@ -128,3 +133,19 @@ test("styles quotes, rules, images and task lists", async ({ page }) => {
   // A checkbox stands in for the bullet.
   expect(await css(reply, "li:has(> input)", "list-style-type")).toBe("none");
 });
+
+test("loads local and web reply images", async ({ page }) => {
+  await page.route("https://images.example.com/code.gif", (route) =>
+    route.fulfill({ contentType: "image/gif", body: Buffer.from(PIXEL.split(",")[1]!, "base64") }),
+  );
+  const reply = await openReply(
+    page,
+    "![Local QR](/tmp/code.png)\n\n![Web QR](https://images.example.com/code.gif)",
+    PIXEL,
+  );
+  for (const name of ["Local QR", "Web QR"]) {
+    await expect
+      .poll(() => reply.getByRole("img", { name }).evaluate((img: HTMLImageElement) => img.naturalWidth))
+      .toBe(1);
+  }
+});

src-tauri/tauri.conf.json

@@ -26,7 +26,7 @@
       }
     ],
     "security": {
-      "csp": "default-src 'self'; img-src 'self' asset: data: blob: https://avatars.githubusercontent.com; style-src 'self' 'unsafe-inline'; font-src 'self' data:; connect-src ipc: http://ipc.localhost; script-src 'self' 'wasm-unsafe-eval'; worker-src 'self' blob:; frame-src http://127.0.0.1:* openidea-plugin: http://openidea-plugin.localhost"
+      "csp": "default-src 'self'; img-src 'self' asset: data: blob: http: https:; style-src 'self' 'unsafe-inline'; font-src 'self' data:; connect-src ipc: http://ipc.localhost; script-src 'self' 'wasm-unsafe-eval'; worker-src 'self' blob:; frame-src http://127.0.0.1:* openidea-plugin: http://openidea-plugin.localhost"
     }
   },
   "bundle": {

src/features/agents/Markdown.test.tsx

@@ -1,15 +1,23 @@
-import { fireEvent, render, screen } from "@testing-library/react";
-import { beforeEach, expect, it, vi } from "vitest";
+import { act, fireEvent, render, screen, waitFor } from "@testing-library/react";
+import { afterEach, beforeEach, expect, it, vi } from "vitest";
+import { clearMocks, mockIPC } from "@tauri-apps/api/mocks";
 import { Markdown } from "./Markdown";
 
 const { openFile, openExternalUrl } = vi.hoisted(() => ({
   openFile: vi.fn().mockResolvedValue(undefined),
   openExternalUrl: vi.fn().mockResolvedValue(undefined),
 }));
+const readImage = vi.fn<(path: string) => Promise<string>>();
+const PIXEL = "data:image/gif;base64,R0lGODlhAQABAIAAAAAAAP///yH5BAEAAAAALAAAAAABAAEAAAIBRAA7";
 vi.mock("@/workbench/contributions", () => ({ editorService: () => ({ openFile }) }));
 vi.mock("@/platform/openUrl", () => ({ openExternalUrl }));
 
-beforeEach(() => vi.clearAllMocks());
+beforeEach(() => {
+  vi.clearAllMocks();
+  readImage.mockReset().mockResolvedValue(PIXEL);
+  mockIPC((cmd, args) => (cmd === "fs_read_image" ? readImage((args as { path: string }).path) : null));
+});
+afterEach(clearMocks);
 
 it.each([
   "<style>body{outline:7px solid red}</style>",
@@ -71,3 +79,90 @@ it.each(["#section", "//example.com/image.jpg", "javascript:alert(1)", "/tmp/%ZZ
     expect(openExternalUrl).not.toHaveBeenCalled();
   },
 );
+
+it.each([
+  ["/Users/thomas/vmos-shots/tiktok-home.jpg", "/repo", "/Users/thomas/vmos-shots/tiktok-home.jpg"],
+  ["images/home.jpg", "/repo", "/repo/images/home.jpg"],
+  ["/tmp/截图%20one.jpg", "/repo", "/tmp/截图 one.jpg"],
+  ["file:///tmp/auth%20code.png", "/repo", "/tmp/auth code.png"],
+  ["C:/shots/code.png", "C:/repo", "C:/shots/code.png"],
+  ["file:///C:/shots/code.png", "C:/repo", "C:/shots/code.png"],
+])("loads a reply image from the project host: %s", async (src, cwd, path) => {
+  render(<Markdown text={`![QR code](${src})`} cwd={cwd} />);
+  await waitFor(() => expect(readImage).toHaveBeenCalledWith(path));
+  await waitFor(() => expect(screen.getByRole("img", { name: "QR code" })).toHaveAttribute("src", PIXEL));
+});
+
+it("loads raw HTML images without allowing alternate sources or forged paths", async () => {
+  const { container } = render(
+    <Markdown
+      text={
+        '<picture><source srcset="/wrong.png"><img alt="QR" src="/tmp/code.png" srcset="/wrong.png 2x" data-image-src="/forged.png" onerror="alert(1)"></picture>'
+      }
+      cwd="/repo"
+    />,
+  );
+  await waitFor(() => expect(screen.getByRole("img")).toHaveAttribute("src", PIXEL));
+  expect(readImage).toHaveBeenCalledExactlyOnceWith("/tmp/code.png");
+  expect(container.querySelector("source, [srcset], [onerror]")).toBeNull();
+});
+
+it.each([PIXEL, "https://example.com/code.png", "http://localhost:9000/code.png"])(
+  "keeps an inline or web image: %s",
+  (src) => {
+    render(<Markdown text={`![QR](${src})`} cwd="/repo" />);
+    expect(screen.getByRole("img")).toHaveAttribute("src", src);
+    expect(screen.getByRole("img")).toHaveAttribute("referrerpolicy", "no-referrer");
+    expect(readImage).not.toHaveBeenCalled();
+  },
+);
+
+it("reuses image reads while a reply streams, but reads again for another message", async () => {
+  const { rerender, unmount } = render(<Markdown text="![QR](/tmp/stream.png)" cwd="/repo" />);
+  await waitFor(() => expect(screen.getByRole("img")).toHaveAttribute("src", PIXEL));
+  rerender(<Markdown text="![QR](/tmp/stream.png)\n\nScan to authorize." cwd="/repo" />);
+  await waitFor(() => expect(screen.getByRole("img")).toHaveAttribute("src", PIXEL));
+  expect(readImage).toHaveBeenCalledTimes(1);
+  unmount();
+  render(<Markdown text="![QR](/tmp/stream.png)" cwd="/repo" />);
+  await waitFor(() => expect(readImage).toHaveBeenCalledTimes(2));
+});
+
+it("reports read failure and retries on the next render", async () => {
+  readImage.mockRejectedValueOnce({ kind: "io", message: "missing", data: null });
+  const { rerender } = render(<Markdown text="![QR](/tmp/missing.png)" cwd="/repo" />);
+  await waitFor(() => expect(screen.getByRole("img")).toHaveAttribute("alt", "Unable to load image: /tmp/missing.png"));
+  expect(screen.getByRole("img")).not.toHaveAttribute("src");
+  rerender(<Markdown text="![QR](/tmp/missing.png)\n\nNow saved." cwd="/repo" />);
+  await waitFor(() => expect(screen.getByRole("img")).toHaveAttribute("src", PIXEL));
+});
+
+it("does not update detached images after a streamed reply replaces them", async () => {
+  let resolve!: (data: string) => void;
+  readImage.mockReturnValueOnce(
+    new Promise((done) => {
+      resolve = done;
+    }),
+  );
+  const { rerender } = render(<Markdown text="![QR](/tmp/pending.png)" cwd="/repo" />);
+  const oldImage = screen.getByRole("img");
+  rerender(<Markdown text="Image removed." cwd="/repo" />);
+  await act(async () => resolve(PIXEL));
+  expect(oldImage).not.toHaveAttribute("src");
+});
+
+it.each(["javascript:alert(1)", "data:text/html,test", "/tmp/%ZZ.png", "/tmp/code%00.png"])(
+  "does not load unsafe or malformed image sources: %s",
+  (src) => {
+    render(<Markdown text={`<img alt="QR" src="${src}">`} cwd="/repo" />);
+    expect(screen.getByRole("img")).not.toHaveAttribute("src");
+    expect(readImage).not.toHaveBeenCalled();
+  },
+);
+
+it("reports network or decode failures without leaving a broken source", () => {
+  render(<Markdown text="![QR](https://example.com/gone.png)" cwd="/repo" />);
+  fireEvent.error(screen.getByRole("img"));
+  expect(screen.getByRole("img")).not.toHaveAttribute("src");
+  expect(screen.getByRole("img")).toHaveAttribute("alt", "Unable to load image: https://example.com/gone.png");
+});

src/features/agents/Markdown.tsx

@@ -13,6 +13,7 @@ import { openExternalUrl } from "@/platform/openUrl";
 import { copyText } from "@/ui";
 import { editorService } from "@/workbench/contributions";
 import { fileHref, pathLink, type PathLink } from "./conversation";
+import { loadImages } from "./markdownImages";
 import styles from "./Session.module.css";
 
 const marked = new Marked({ gfm: true, breaks: true });
@@ -21,6 +22,26 @@ const marked = new Marked({ gfm: true, breaks: true });
 const rendered = new Map<string, string>();
 const RENDERED_MAX = 500;
 
+let purifier: ReturnType<typeof DOMPurify> | null = null;
+
+function purify(): ReturnType<typeof DOMPurify> {
+  if (purifier) return purifier;
+  const p = DOMPurify(window);
+  const sources = new WeakMap<Node, string>();
+  // Preserve image-only file URLs without relaxing URI sanitizing for links or other tags.
+  p.addHook("beforeSanitizeAttributes", (node) => {
+    if (node.tagName !== "IMG") return;
+    sources.set(node, node.getAttribute("src") ?? "");
+    node.removeAttribute("src");
+  });
+  p.addHook("afterSanitizeAttributes", (node) => {
+    const src = sources.get(node);
+    if (src) node.setAttribute("data-image-src", src);
+  });
+  purifier = p;
+  return p;
+}
+
 export function markdownHtml(text: string): string {
   let html = rendered.get(text);
   if (html === undefined) {
@@ -35,11 +56,11 @@ function renderHtml(text: string): string {
   const html = marked.parse(text, { async: false }) as string;
   // Replies can quote repository content. Styles would reach the whole IDE, and the click
   // handlers below trust `data-copy` / `data-path`, which only `enhance` may set.
-  const root = DOMPurify.sanitize(html, {
+  const root = purify().sanitize(html, {
     USE_PROFILES: { html: true },
     RETURN_DOM: true,
-    FORBID_TAGS: ["style", "form"],
-    FORBID_ATTR: ["style"],
+    FORBID_TAGS: ["style", "form", "source"],
+    FORBID_ATTR: ["style", "srcset"],
     ALLOW_DATA_ATTR: false,
   }) as HTMLElement;
   root.querySelectorAll("li").forEach(wrapLooseText);
@@ -92,8 +113,9 @@ export function openPathLink(link: PathLink) {
 }
 
 /** Highlighting, Copy buttons and file links over freshly rendered `root`. */
-function enhance(root: HTMLElement, cwd: string): () => void {
+function enhance(root: HTMLElement, cwd: string, images: Map<string, Promise<string>>): () => void {
   let alive = true;
+  const disposeImages = loadImages(root, cwd, images);
   root.querySelectorAll<HTMLElement>("pre > code").forEach((code) => {
     const pre = code.parentElement!;
     const lang = [...code.classList].find((c) => c.startsWith("language-"))?.slice("language-".length) ?? null;
@@ -125,13 +147,15 @@ function enhance(root: HTMLElement, cwd: string): () => void {
   });
   return () => {
     alive = false;
+    disposeImages();
   };
 }
 
 export function Markdown({ text, cwd }: { text: string; cwd: string }) {
   const html = useMemo(() => markdownHtml(text), [text]);
   const ref = useRef<HTMLDivElement>(null);
-  useLayoutEffect(() => (ref.current ? enhance(ref.current, cwd) : undefined), [html, cwd]);
+  const images = useRef(new Map<string, Promise<string>>());
+  useLayoutEffect(() => (ref.current ? enhance(ref.current, cwd, images.current) : undefined), [html, cwd]);
   return (
     <div
       ref={ref}

src/features/agents/markdownImages.test.ts

@@ -0,0 +1,37 @@
+import { expect, it } from "vitest";
+import { imageTarget } from "./markdownImages";
+
+it.each([
+  ["images/code.png", "/repo", "/repo/images/code.png"],
+  ["../code.png?v=2#preview", "/repo", "/repo/../code.png"],
+  ["/tmp/截图%20one%23.png", "/repo", "/tmp/截图 one#.png"],
+  ["file://localhost/tmp/code.png", "/repo", "/tmp/code.png"],
+  ["C:/shots/code.png", "/repo", "C:/shots/code.png"],
+  [String.raw`C:\shots\code.png`, "C:/repo", "C:/shots/code.png"],
+  [String.raw`images\code.png`, String.raw`C:\repo`, "C:/repo/images/code.png"],
+  [String.raw`\\server\share\code.png`, "C:/repo", "//server/share/code.png"],
+  ["file://server/share/code.png", "C:/repo", "//server/share/code.png"],
+  ["file:///C:/shots/code%20one.png", "C:/repo", "C:/shots/code one.png"],
+])("resolves host image paths: %s", (src, cwd, path) => {
+  expect(imageTarget(src, cwd)).toEqual({ path });
+});
+
+it.each([
+  "",
+  "#section",
+  "javascript:alert(1)",
+  "JaVaScRiPt:alert(1)",
+  "data:text/html,test",
+  "ftp://example.com/x.png",
+  "/tmp/%ZZ.png",
+  "/tmp/code%00.png",
+  "file:///tmp/%ZZ.png",
+  "https://",
+  "https://example.com/\ncode.png",
+])("rejects unsupported or malformed sources: %s", (src) => {
+  expect(imageTarget(src, "/repo")).toBeNull();
+});
+
+it("treats protocol-relative URLs as HTTPS, rather than host files", () => {
+  expect(imageTarget("//example.com/code.png", "/repo")).toEqual({ url: "https://example.com/code.png" });
+});

src/features/agents/markdownImages.ts

@@ -0,0 +1,75 @@
+import { call, commands } from "@/platform/ipc";
+import { ipcPath } from "@/platform/paths";
+import { tx } from "@/platform/i18n";
+
+const hasControl = (text: string) => [...text].some((c) => c.charCodeAt(0) < 32 || c.charCodeAt(0) === 127);
+
+type ImageTarget = { url: string } | { path: string };
+
+/** Images name host files, not paths relative to the webview's own origin. */
+export function imageTarget(source: string, cwd: string): ImageTarget | null {
+  const src = source.trim();
+  if (!src || hasControl(src)) return null;
+  if (/^data:image\//i.test(src)) return { url: src };
+  try {
+    if (/^https?:\/\//i.test(src) || src.startsWith("//")) {
+      const url = src.startsWith("//") ? `https:${src}` : src;
+      new URL(url);
+      return { url };
+    }
+    let path: string;
+    if (/^file:/i.test(src)) {
+      const url = new URL(src);
+      path = decodeURIComponent(url.pathname);
+      if (/^\/[a-z]:\//i.test(path)) path = path.slice(1);
+      if (url.hostname && url.hostname !== "localhost") path = `//${url.hostname}${path}`;
+    } else {
+      if (/^[a-z][a-z\d+.-]*:/i.test(src) && !/^[a-z]:[\\/]/i.test(src)) return null;
+      path = ipcPath(decodeURIComponent(src.split(/[?#]/, 1)[0]!));
+    }
+    if (!path || hasControl(path)) return null;
+    if (path.startsWith("/") || /^[a-z]:\//i.test(path)) return { path };
+    const base = ipcPath(cwd).replace(/\/$/, "");
+    if (/^[a-z]:\//i.test(base) || base.startsWith("//")) path = path.replace(/\\/g, "/");
+    return { path: `${base}/${path}` };
+  } catch {
+    return null;
+  }
+}
+
+/** Cache belongs to one message: streaming reuses reads, later replies can show a new QR at the same path. */
+export function loadImages(root: HTMLElement, cwd: string, cache: Map<string, Promise<string>>): () => void {
+  let alive = true;
+  const images = [...root.querySelectorAll<HTMLImageElement>("img")];
+  for (const img of images) {
+    const source = img.dataset.imageSrc ?? "";
+    const fail = () => {
+      if (!alive) return;
+      img.onerror = null;
+      img.removeAttribute("src");
+      img.alt = tx`Unable to load image: ${source}`;
+    };
+    img.setAttribute("referrerpolicy", "no-referrer");
+    img.onerror = fail;
+    const target = imageTarget(source, cwd);
+    if (!target) {
+      fail();
+    } else if ("url" in target) {
+      img.src = target.url;
+    } else {
+      let pending = cache.get(target.path);
+      if (!pending) {
+        pending = call(commands.fsReadImage(target.path));
+        cache.set(target.path, pending);
+        void pending.catch(() => cache.delete(target.path));
+      }
+      void pending.then((url) => {
+        if (alive) img.src = url;
+      }, fail);
+    }
+  }
+  return () => {
+    alive = false;
+    for (const img of images) img.onerror = null;
+  };
+}

src/platform/locales/text/en.json

@@ -36,5 +36,6 @@
   "For web requests: the XDEBUG_SESSION cookie or query parameter; point xdebug.client_host at this machine when PHP runs in a container": "For web requests: the XDEBUG_SESSION cookie or query parameter; point xdebug.client_host at this machine when PHP runs in a container",
   "e.g. /var/www/html": "e.g. /var/www/html",
   "Listening for Xdebug on port {port}": "Listening for Xdebug on port {port}",
-  "Port {port} is already in use: another IDE or debug session may be listening on it. Stop it or choose another port.": "Port {port} is already in use: another IDE or debug session may be listening on it. Stop it or choose another port."
+  "Port {port} is already in use: another IDE or debug session may be listening on it. Stop it or choose another port.": "Port {port} is already in use: another IDE or debug session may be listening on it. Stop it or choose another port.",
+  "Unable to load image: {path}": "Unable to load image: {path}"
 }

src/platform/locales/text/zh.json

@@ -6376,5 +6376,6 @@
   "Force-with-lease push {source} to {target} was rejected. Force-with-lease push prevents overriding remote changes that haven't been fetched to the local repository. Fetch recent changes to verify that they can be discarded safely and repeat the push operation.": "将 {source} 以 force-with-lease 方式推送到 {target} 被拒绝。force-with-lease 推送可防止覆盖尚未获取到本地仓库的远程更改。请先获取最新更改，确认可以安全丢弃后再重新推送。",
   "Force push anyway": "仍然强制推送",
   "Stage and Retry": "暂存并重试",
-  "Show Files": "显示文件"
+  "Show Files": "显示文件",
+  "Unable to load image: {path}": "无法加载图片：{path}"
 }