diff --git a/apps/code/src/renderer/components/ui/SafeImagePreview.tsx b/apps/code/src/renderer/components/ui/SafeImagePreview.tsx new file mode 100644 index 000000000..cbe6fbd63 --- /dev/null +++ b/apps/code/src/renderer/components/ui/SafeImagePreview.tsx @@ -0,0 +1,63 @@ +import { Flex, Text } from "@radix-ui/themes"; +import { + buildImageDataUrl, + isAllowedImageMimeType, + MAX_IMAGE_BASE64_LENGTH, +} from "@shared/utils/imageDataUrl"; +import { useState } from "react"; + +interface SafeImagePreviewProps { + /** Base64-encoded image data (no data URL prefix). */ + base64: string; + mimeType: string; + alt?: string; + className?: string; + /** Rendered when the image fails to decode or has a disallowed mime type. */ + fallback?: React.ReactNode; +} + +function DefaultFallback() { + return ( + + Unable to render image preview + + ); +} + +export function SafeImagePreview({ + base64, + mimeType, + alt, + className, + fallback, +}: SafeImagePreviewProps) { + const [hasError, setHasError] = useState(false); + const [lastSource, setLastSource] = useState({ base64, mimeType }); + + if (lastSource.base64 !== base64 || lastSource.mimeType !== mimeType) { + setLastSource({ base64, mimeType }); + setHasError(false); + } + + const isPayloadValid = + base64.length > 0 && + base64.length <= MAX_IMAGE_BASE64_LENGTH && + isAllowedImageMimeType(mimeType); + + if (!isPayloadValid || hasError) { + return <>{fallback ?? }; + } + + return ( + {alt setHasError(true)} + /> + ); +} diff --git a/apps/code/src/renderer/features/code-editor/components/CodeEditorPanel.tsx b/apps/code/src/renderer/features/code-editor/components/CodeEditorPanel.tsx index 45bbb5090..dccd768e9 100644 --- a/apps/code/src/renderer/features/code-editor/components/CodeEditorPanel.tsx +++ b/apps/code/src/renderer/features/code-editor/components/CodeEditorPanel.tsx @@ -1,4 +1,5 @@ import { PanelMessage } from "@components/ui/PanelMessage"; +import { SafeImagePreview } from "@components/ui/SafeImagePreview"; import { Tooltip } from "@components/ui/Tooltip"; import { CodeMirrorEditor } from "@features/code-editor/components/CodeMirrorEditor"; import { EnrichmentPopover } from "@features/code-editor/components/EnrichmentPopover"; @@ -16,6 +17,7 @@ import { Box, Flex, IconButton, Text } from "@radix-ui/themes"; import { trpcClient, useTRPC } from "@renderer/trpc/client"; import { getImageMimeType, isImageFile } from "@shared/constants/image"; import type { Task } from "@shared/types"; +import { parseImageDataUrl } from "@shared/utils/imageDataUrl"; import { useQuery } from "@tanstack/react-query"; import { useCallback, useMemo, useState } from "react"; @@ -29,6 +31,40 @@ interface CodeEditorPanelProps { absolutePath: string; } +function FilePanelImagePreview({ + base64, + mimeType, + filePath, + absolutePath, +}: { + base64: string; + mimeType: string; + filePath: string; + absolutePath: string; +}) { + return ( + + + Failed to render image + + } + /> + + ); +} + export function CodeEditorPanel({ taskId, task: _task, @@ -128,6 +164,12 @@ export function CodeEditorPanel({ content: isImage ? null : fileContent, }); + const dataUrlImage = useMemo( + () => + isImage || fileContent == null ? null : parseImageDataUrl(fileContent), + [isImage, fileContent], + ); + if (isImage) { if (isCloudRun) { return ( @@ -144,21 +186,13 @@ export function CodeEditorPanel({ Failed to load image ); } - const mimeType = getImageMimeType(absolutePath); return ( - - {filePath} - + ); } @@ -192,6 +226,17 @@ export function CodeEditorPanel({ return File is empty; } + if (dataUrlImage) { + return ( + + ); + } + if (isMarkdown) { const handleCopySource = () => { navigator.clipboard.writeText(fileContent); diff --git a/apps/code/src/renderer/features/sessions/components/session-update/CodePreview.tsx b/apps/code/src/renderer/features/sessions/components/session-update/CodePreview.tsx index 0efe74dca..7f7841d28 100644 --- a/apps/code/src/renderer/features/sessions/components/session-update/CodePreview.tsx +++ b/apps/code/src/renderer/features/sessions/components/session-update/CodePreview.tsx @@ -1,6 +1,8 @@ import { EditorView } from "@codemirror/view"; +import { SafeImagePreview } from "@components/ui/SafeImagePreview"; import { MultiFileDiff } from "@pierre/diffs/react"; import { Code } from "@radix-ui/themes"; +import { parseImageDataUrl } from "@shared/utils/imageDataUrl"; import { useThemeStore } from "@stores/themeStore"; import { compactHomePath } from "@utils/path"; import { useEffect, useMemo, useRef } from "react"; @@ -31,6 +33,10 @@ export function CodePreview({ cacheKey, }: CodePreviewProps) { const isDiff = oldContent !== undefined && oldContent !== null; + const imageDataUrl = useMemo( + () => (isDiff ? null : parseImageDataUrl(content)), + [isDiff, content], + ); if (isDiff) { return ( @@ -45,6 +51,18 @@ export function CodePreview({ ); } + if (imageDataUrl) { + return ( + + ); + } + return ( + {showPath && filePath && ( +
+ + {compactHomePath(filePath)} + +
+ )} +
+ +
+ + ); +} + function DiffPreview({ content, filePath, diff --git a/apps/code/src/renderer/features/sessions/components/session-update/ReadToolView.tsx b/apps/code/src/renderer/features/sessions/components/session-update/ReadToolView.tsx index d10b9c07e..779943e6b 100644 --- a/apps/code/src/renderer/features/sessions/components/session-update/ReadToolView.tsx +++ b/apps/code/src/renderer/features/sessions/components/session-update/ReadToolView.tsx @@ -1,3 +1,4 @@ +import { SafeImagePreview } from "@components/ui/SafeImagePreview"; import { FileText } from "@phosphor-icons/react"; import { Box, Flex } from "@radix-ui/themes"; import { useState } from "react"; @@ -5,6 +6,7 @@ import { CodePreview } from "./CodePreview"; import { FileMentionChip } from "./FileMentionChip"; import { ExpandableIcon, + getContentImage, getReadToolContent, StatusIndicators, ToolTitle, @@ -27,9 +29,10 @@ export function ReadToolView({ const filePath = locations?.[0]?.path ?? ""; const startLine = locations?.[0]?.line ?? 0; - const fileContent = getReadToolContent(content); + const imageContent = getContentImage(content); + const fileContent = imageContent ? undefined : getReadToolContent(content); const lineCount = fileContent ? fileContent.split("\n").length : null; - const isExpandable = !!fileContent; + const isExpandable = !!fileContent || !!imageContent; const firstLineNumber = startLine + 1; const handleClick = () => { @@ -53,12 +56,27 @@ export function ReadToolView({ isExpanded={isExpanded} /> - Read{lineCount !== null ? ` ${lineCount} lines in` : ""} + {imageContent + ? "Read image in" + : `Read${lineCount !== null ? ` ${lineCount} lines in` : ""}`} {filePath && } + {isExpanded && imageContent && ( + + + + + + )} + {isExpanded && fileContent && ( diff --git a/apps/code/src/renderer/features/sessions/components/session-update/toolCallUtils.tsx b/apps/code/src/renderer/features/sessions/components/session-update/toolCallUtils.tsx index bc49e457e..eb73b3da4 100644 --- a/apps/code/src/renderer/features/sessions/components/session-update/toolCallUtils.tsx +++ b/apps/code/src/renderer/features/sessions/components/session-update/toolCallUtils.tsx @@ -68,6 +68,26 @@ export function getContentText( return undefined; } +export interface ImageContentData { + base64: string; + mimeType: string; +} + +export function getContentImage( + content: ToolCall["content"], +): ImageContentData | undefined { + if (!content?.length) return undefined; + for (const item of content) { + if (item.type === "content" && item.content.type === "image") { + const { data, mimeType } = item.content; + if (typeof data === "string" && typeof mimeType === "string") { + return { base64: data, mimeType }; + } + } + } + return undefined; +} + export function getReadToolContent( content: ToolCall["content"], ): string | undefined { diff --git a/apps/code/src/shared/constants/image.ts b/apps/code/src/shared/constants/image.ts index a201592ad..b7b2e11fc 100644 --- a/apps/code/src/shared/constants/image.ts +++ b/apps/code/src/shared/constants/image.ts @@ -6,7 +6,6 @@ export const IMAGE_MIME_TYPES: Record = { webp: "image/webp", bmp: "image/bmp", ico: "image/x-icon", - svg: "image/svg+xml", tiff: "image/tiff", tif: "image/tiff", }; diff --git a/apps/code/src/shared/utils/imageDataUrl.test.ts b/apps/code/src/shared/utils/imageDataUrl.test.ts new file mode 100644 index 000000000..f197142da --- /dev/null +++ b/apps/code/src/shared/utils/imageDataUrl.test.ts @@ -0,0 +1,143 @@ +import { describe, expect, it } from "vitest"; +import { + buildImageDataUrl, + isAllowedImageMimeType, + parseImageDataUrl, +} from "./imageDataUrl"; + +const TINY_PNG_BASE64 = + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII="; + +describe("parseImageDataUrl", () => { + it("parses a valid PNG data URL", () => { + const result = parseImageDataUrl( + `data:image/png;base64,${TINY_PNG_BASE64}`, + ); + expect(result).toEqual({ + mimeType: "image/png", + base64: TINY_PNG_BASE64, + }); + }); + + it.each([ + ["image/jpeg"], + ["image/webp"], + ["image/gif"], + ["image/bmp"], + ["image/avif"], + ["image/tiff"], + ["image/x-icon"], + ])("accepts allowed mime type %s", (mimeType) => { + const result = parseImageDataUrl( + `data:${mimeType};base64,${TINY_PNG_BASE64}`, + ); + expect(result).not.toBeNull(); + expect(result?.mimeType).toBe(mimeType); + }); + + it("rejects SVG data URLs to prevent script execution", () => { + expect( + parseImageDataUrl(`data:image/svg+xml;base64,${TINY_PNG_BASE64}`), + ).toBeNull(); + }); + + it.each([ + ["text/html"], + ["application/javascript"], + ["application/octet-stream"], + ["text/plain"], + ])("rejects non-image mime type %s", (mimeType) => { + expect( + parseImageDataUrl(`data:${mimeType};base64,${TINY_PNG_BASE64}`), + ).toBeNull(); + }); + + it("rejects non-base64 data URLs", () => { + expect(parseImageDataUrl("data:image/png,not-base64")).toBeNull(); + }); + + it.each([ + ["empty string", ""], + ["plain text", "hello world"], + ["http URL", "https://example.com/image.png"], + ["truncated data prefix", "data"], + ["missing payload separator", "data:image/png;base64"], + ["empty payload", "data:image/png;base64,"], + ["bare prefix", "data:"], + ])("rejects non-data-URL or malformed input: %s", (_label, value) => { + expect(parseImageDataUrl(value)).toBeNull(); + }); + + it("rejects extremely large payloads", () => { + const huge = "A".repeat(30 * 1024 * 1024); + expect(parseImageDataUrl(`data:image/png;base64,${huge}`)).toBeNull(); + }); + + it("trims surrounding whitespace before parsing", () => { + const result = parseImageDataUrl( + `\n data:image/png;base64,${TINY_PNG_BASE64} \n`, + ); + expect(result?.mimeType).toBe("image/png"); + }); + + it("tolerates long leading-whitespace prefixes", () => { + const padding = " ".repeat(256); + const result = parseImageDataUrl( + `${padding}data:image/png;base64,${TINY_PNG_BASE64}`, + ); + expect(result?.mimeType).toBe("image/png"); + }); + + it("strips whitespace inside base64 payload", () => { + const withNewlines = TINY_PNG_BASE64.match(/.{1,40}/g)?.join("\n") ?? ""; + const result = parseImageDataUrl(`data:image/png;base64,${withNewlines}`); + expect(result?.base64).toBe(TINY_PNG_BASE64); + }); + + it("ignores additional parameters before the base64 marker", () => { + const result = parseImageDataUrl( + `data:image/png;charset=utf-8;base64,${TINY_PNG_BASE64}`, + ); + expect(result?.mimeType).toBe("image/png"); + }); + + it("normalises mime type casing", () => { + const result = parseImageDataUrl( + `data:IMAGE/PNG;base64,${TINY_PNG_BASE64}`, + ); + expect(result?.mimeType).toBe("image/png"); + }); + + it.each([[null], [undefined], [123], [{}]])( + "handles non-string input safely: %p", + (value) => { + expect(parseImageDataUrl(value as unknown as string)).toBeNull(); + }, + ); +}); + +describe("isAllowedImageMimeType", () => { + it.each([["image/png"], ["IMAGE/JPEG"], ["image/webp"], ["image/gif"]])( + "accepts %s", + (mimeType) => { + expect(isAllowedImageMimeType(mimeType)).toBe(true); + }, + ); + + it.each([ + ["image/svg+xml"], + ["text/html"], + ["application/javascript"], + ["text/plain"], + ])("rejects %s", (mimeType) => { + expect(isAllowedImageMimeType(mimeType)).toBe(false); + }); +}); + +describe("buildImageDataUrl", () => { + it("builds a data URL from parts", () => { + expect(buildImageDataUrl("image/png", "abc")).toBe( + "data:image/png;base64,abc", + ); + }); +}); diff --git a/apps/code/src/shared/utils/imageDataUrl.ts b/apps/code/src/shared/utils/imageDataUrl.ts new file mode 100644 index 000000000..7ce852265 --- /dev/null +++ b/apps/code/src/shared/utils/imageDataUrl.ts @@ -0,0 +1,52 @@ +const ALLOWED_IMAGE_MIME_TYPES = new Set([ + "image/png", + "image/jpeg", + "image/gif", + "image/webp", + "image/bmp", + "image/x-icon", + "image/vnd.microsoft.icon", + "image/tiff", + "image/avif", +]); + +const DATA_URL_PATTERN = + /^data:([a-zA-Z]+\/[a-zA-Z0-9.+-]+)(?:;[a-zA-Z0-9-]+=[^;,]+)*;base64,([A-Za-z0-9+/=\s]+)$/; + +const MAX_DATA_URL_LENGTH = 20 * 1024 * 1024; +export const MAX_IMAGE_BASE64_LENGTH = 15 * 1024 * 1024; + +export interface ParsedImageDataUrl { + mimeType: string; + base64: string; +} + +export function parseImageDataUrl(value: string): ParsedImageDataUrl | null { + if (typeof value !== "string" || value.length === 0) return null; + if (value.length > MAX_DATA_URL_LENGTH) return null; + if (!/^\s{0,1024}data:/.test(value)) return null; + + const trimmed = value.trim(); + if (trimmed.length === 0) return null; + + const match = DATA_URL_PATTERN.exec(trimmed); + if (!match) return null; + + const mimeType = match[1].toLowerCase(); + if (!ALLOWED_IMAGE_MIME_TYPES.has(mimeType)) return null; + + const base64 = match[2].replace(/\s+/g, ""); + if (base64.length === 0 || base64.length > MAX_IMAGE_BASE64_LENGTH) { + return null; + } + + return { mimeType, base64 }; +} + +export function isAllowedImageMimeType(mimeType: string): boolean { + return ALLOWED_IMAGE_MIME_TYPES.has(mimeType.toLowerCase()); +} + +export function buildImageDataUrl(mimeType: string, base64: string): string { + return `data:${mimeType};base64,${base64}`; +}