feat: optimize memory usage for image handling in webview (#7556)

* feat: optimize memory usage for image handling in webview

- Replace base64 image data with webview URIs to reduce memory footprint
- Add proper resource roots to webview for workspace file access
- Implement convertToWebviewUri method for safe file-to-URI conversion
- Update ImageViewer to handle both webview URIs and file paths separately
- Add image message type for proper image rendering in chat
- Improve error handling and display for failed image loads
- Add comprehensive tests for ImageViewer component
- Format display paths as relative for better readability

This change significantly reduces memory usage by avoiding base64 encoding
of images and instead using VSCode's webview URI system for direct file
access. Images are now loaded on-demand from disk rather than being held
in memory as base64 strings.

* fix: address PR review comments

- Use safeJsonParse instead of JSON.parse in ChatRow.tsx
- Add type definition for parsed image info
- Add more specific error types in ClineProvider.ts
- Add comprehensive JSDoc comments to ImageBlock.tsx
- Improve error handling and type safety

* fix: address MrUbens' review comments

- Remove hardcoded 'rc1' pattern in formatDisplayPath, use generic workspace detection
- Internationalize 'No image data' text using i18n system

* chore: remove useless comment

* chore(i18n): add image.noData to all locales to fix translation check

* test: update ImageViewer.spec to align with i18n key and flexible path formatting

Author: daniel-lxs

src/core/tools/generateImageTool.ts

@@ -8,7 +8,6 @@ import { fileExistsAtPath } from "../../utils/fs"
 import { getReadablePath } from "../../utils/path"
 import { isPathOutsideWorkspace } from "../../utils/pathUtils"
 import { EXPERIMENT_IDS, experiments } from "../../shared/experiments"
-import { safeWriteJson } from "../../utils/safeWriteJson"
 import { OpenRouterHandler } from "../../api/providers/openrouter"
 
 // Hardcoded list of image generation models for now
@@ -237,12 +236,18 @@ export async function generateImageTool(
 
 			cline.didEditFile = true
 
-			// Display the generated image in the chat using a text message with the image
-			await cline.say("text", getReadablePath(cline.cwd, finalPath), [result.imageData])
-
 			// Record successful tool usage
 			cline.recordToolUsage("generate_image")
 
+			// Get the webview URI for the image
+			const provider = cline.providerRef.deref()
+			const fullImagePath = path.join(cline.cwd, finalPath)
+
+			// Convert to webview URI if provider is available
+			const imageUri = provider?.convertToWebviewUri?.(fullImagePath) ?? vscode.Uri.file(fullImagePath).toString()
+
+			// Send the image with the webview URI
+			await cline.say("image", JSON.stringify({ imageUri, imagePath: fullImagePath }))
 			pushToolResult(formatResponse.toolResult(getReadablePath(cline.cwd, finalPath)))
 
 			return

src/core/webview/ClineProvider.ts

@@ -604,9 +604,17 @@ export class ClineProvider
 			setTtsSpeed(ttsSpeed ?? 1)
 		})
 
+		// Set up webview options with proper resource roots
+		const resourceRoots = [this.contextProxy.extensionUri]
+
+		// Add workspace folders to allow access to workspace files
+		if (vscode.workspace.workspaceFolders) {
+			resourceRoots.push(...vscode.workspace.workspaceFolders.map((folder) => folder.uri))
+		}
+
 		webviewView.webview.options = {
 			enableScripts: true,
-			localResourceRoots: [this.contextProxy.extensionUri],
+			localResourceRoots: resourceRoots,
 		}
 
 		webviewView.webview.html =
@@ -2500,4 +2508,40 @@ public async clearTask(): Promise<void> {
 	public get cwd() {
 		return getWorkspacePath()
 	}
+
+	/**
+	 * Convert a file path to a webview-accessible URI
+	 * This method safely converts file paths to URIs that can be loaded in the webview
+	 *
+	 * @param filePath - The absolute file path to convert
+	 * @returns The webview URI string, or the original file URI if conversion fails
+	 * @throws {Error} When webview is not available
+	 * @throws {TypeError} When file path is invalid
+	 */
+	public convertToWebviewUri(filePath: string): string {
+		try {
+			const fileUri = vscode.Uri.file(filePath)
+
+			// Check if we have a webview available
+			if (this.view?.webview) {
+				const webviewUri = this.view.webview.asWebviewUri(fileUri)
+				return webviewUri.toString()
+			}
+
+			// Specific error for no webview available
+			const error = new Error("No webview available for URI conversion")
+			console.error(error.message)
+			// Fallback to file URI if no webview available
+			return fileUri.toString()
+		} catch (error) {
+			// More specific error handling
+			if (error instanceof TypeError) {
+				console.error("Invalid file path provided for URI conversion:", error)
+			} else {
+				console.error("Failed to convert to webview URI:", error)
+			}
+			// Return file URI as fallback
+			return vscode.Uri.file(filePath).toString()
+		}
+	}
 }

webview-ui/src/components/chat/ChatRow.tsx

@@ -1159,6 +1159,17 @@ export const ChatRowContent = ({
 					return <CodebaseSearchResultsDisplay results={results} />
 				case "user_edit_todos":
 					return <UpdateTodoListToolBlock userEdited onChange={() => {}} />
+				case "image":
+					// Parse the JSON to get imageUri and imagePath
+					const imageInfo = safeJsonParse<{ imageUri: string; imagePath: string }>(message.text || "{}")
+					if (!imageInfo) {
+						return null
+					}
+					return (
+						<div style={{ marginTop: "10px" }}>
+							<ImageBlock imageUri={imageInfo.imageUri} imagePath={imageInfo.imagePath} />
+						</div>
+					)
 				default:
 					return (
 						<>

webview-ui/src/components/common/ImageBlock.tsx

@@ -1,15 +1,66 @@
 import React from "react"
 import { ImageViewer } from "./ImageViewer"
 
+/**
+ * Props for the ImageBlock component
+ */
 interface ImageBlockProps {
-	imageData: string
+	/**
+	 * The webview-accessible URI for rendering the image.
+	 * This is the preferred format for new image generation tools.
+	 * Should be a URI that can be directly loaded in the webview context.
+	 */
+	imageUri?: string
+
+	/**
+	 * The actual file path for display purposes and file operations.
+	 * Used to show the path to the user and for opening the file in the editor.
+	 * This is typically an absolute or relative path to the image file.
+	 */
+	imagePath?: string
+
+	/**
+	 * Base64 data or regular URL for backward compatibility.
+	 * @deprecated Use imageUri instead for new implementations.
+	 * This is maintained for compatibility with Mermaid diagrams and legacy code.
+	 */
+	imageData?: string
+
+	/**
+	 * Optional path for Mermaid diagrams.
+	 * @deprecated Use imagePath instead for new implementations.
+	 * This is maintained for backward compatibility with existing Mermaid diagram rendering.
+	 */
 	path?: string
 }
 
-export default function ImageBlock({ imageData, path }: ImageBlockProps) {
+export default function ImageBlock({ imageUri, imagePath, imageData, path }: ImageBlockProps) {
+	// Determine which props to use based on what's provided
+	let finalImageUri: string
+	let finalImagePath: string | undefined
+
+	if (imageUri) {
+		// New format: explicit imageUri and imagePath
+		finalImageUri = imageUri
+		finalImagePath = imagePath
+	} else if (imageData) {
+		// Legacy format: use imageData as direct URI (for Mermaid diagrams)
+		finalImageUri = imageData
+		finalImagePath = path
+	} else {
+		// No valid image data provided
+		console.error("ImageBlock: No valid image data provided")
+		return null
+	}
+
 	return (
 		<div className="my-2">
-			<ImageViewer imageData={imageData} path={path} alt="AI Generated Image" showControls={true} />
+			<ImageViewer
+				imageUri={finalImageUri}
+				imagePath={finalImagePath}
+				alt="AI Generated Image"
+				showControls={true}
+			/>
 		</div>
 	)
 }

webview-ui/src/components/common/ImageViewer.tsx

@@ -13,17 +13,17 @@ const MIN_ZOOM = 0.5
 const MAX_ZOOM = 20
 
 export interface ImageViewerProps {
-	imageData: string // base64 data URL or regular URL
+	imageUri: string // The URI to use for rendering (webview URI, base64, or regular URL)
+	imagePath?: string // The actual file path for display and opening
 	alt?: string
-	path?: string
 	showControls?: boolean
 	className?: string
 }
 
 export function ImageViewer({
-	imageData,
+	imageUri,
+	imagePath,
 	alt = "Generated image",
-	path,
 	showControls = true,
 	className = "",
 }: ImageViewerProps) {
@@ -33,6 +33,7 @@ export function ImageViewer({
 	const [isHovering, setIsHovering] = useState(false)
 	const [isDragging, setIsDragging] = useState(false)
 	const [dragPosition, setDragPosition] = useState({ x: 0, y: 0 })
+	const [imageError, setImageError] = useState<string | null>(null)
 	const { copyWithFeedback } = useCopyToClipboard()
 	const { t } = useAppTranslation()
 
@@ -53,12 +54,13 @@ export function ImageViewer({
 		e.stopPropagation()
 
 		try {
-			const textToCopy = path || imageData
-			await copyWithFeedback(textToCopy, e)
-
-			// Show feedback
-			setCopyFeedback(true)
-			setTimeout(() => setCopyFeedback(false), 2000)
+			// Copy the file path if available
+			if (imagePath) {
+				await copyWithFeedback(imagePath, e)
+				// Show feedback
+				setCopyFeedback(true)
+				setTimeout(() => setCopyFeedback(false), 2000)
+			}
 		} catch (err) {
 			console.error("Error copying:", err instanceof Error ? err.message : String(err))
 		}
@@ -71,10 +73,10 @@ export function ImageViewer({
 		e.stopPropagation()
 
 		try {
-			// Send message to VSCode to save the image
+			// Request VSCode to save the image
 			vscode.postMessage({
 				type: "saveImage",
-				dataUri: imageData,
+				dataUri: imageUri,
 			})
 		} catch (error) {
 			console.error("Error saving image:", error)
@@ -86,10 +88,21 @@ export function ImageViewer({
 	 */
 	const handleOpenInEditor = (e: React.MouseEvent) => {
 		e.stopPropagation()
-		vscode.postMessage({
-			type: "openImage",
-			text: imageData,
-		})
+		// Use openImage for both file paths and data URIs
+		// The backend will handle both cases appropriately
+		if (imagePath) {
+			// Use the actual file path for opening
+			vscode.postMessage({
+				type: "openImage",
+				text: imagePath,
+			})
+		} else if (imageUri) {
+			// Fallback to opening image URI if no path is available (for Mermaid diagrams)
+			vscode.postMessage({
+				type: "openImage",
+				text: imageUri,
+			})
+		}
 	}
 
 	/**
@@ -129,24 +142,86 @@ export function ImageViewer({
 		setIsHovering(false)
 	}
 
+	const handleImageError = useCallback(() => {
+		setImageError("Failed to load image")
+	}, [])
+
+	const handleImageLoad = useCallback(() => {
+		setImageError(null)
+	}, [])
+
+	/**
+	 * Format the display path for the image
+	 */
+	const formatDisplayPath = (path: string): string => {
+		// If it's already a relative path starting with ./, keep it
+		if (path.startsWith("./")) return path
+		// If it's an absolute path, extract the relative portion
+		// Look for workspace patterns - match the last segment after any directory separator
+		const workspaceMatch = path.match(/\/([^/]+)\/(.+)$/)
+		if (workspaceMatch && workspaceMatch[2]) {
+			// Return relative path from what appears to be the workspace root
+			return `./${workspaceMatch[2]}`
+		}
+		// Otherwise, just get the filename
+		const filename = path.split("/").pop()
+		return filename || path
+	}
+
+	// Handle missing image URI
+	if (!imageUri) {
+		return (
+			<div
+				className={`relative w-full ${className}`}
+				style={{
+					minHeight: "100px",
+					backgroundColor: "var(--vscode-editor-background)",
+					display: "flex",
+					alignItems: "center",
+					justifyContent: "center",
+				}}>
+				<span style={{ color: "var(--vscode-descriptionForeground)" }}>{t("common:image.noData")}</span>
+			</div>
+		)
+	}
+
 	return (
 		<>
 			<div
 				className={`relative w-full ${className}`}
 				onMouseEnter={handleMouseEnter}
 				onMouseLeave={handleMouseLeave}>
-				<img
-					src={imageData}
-					alt={alt}
-					className="w-full h-auto rounded cursor-pointer"
-					onClick={handleOpenInEditor}
-					style={{
-						maxHeight: "400px",
-						objectFit: "contain",
-						backgroundColor: "var(--vscode-editor-background)",
-					}}
-				/>
-				{path && <div className="mt-1 text-xs text-vscode-descriptionForeground">{path}</div>}
+				{imageError ? (
+					<div
+						style={{
+							minHeight: "100px",
+							display: "flex",
+							alignItems: "center",
+							justifyContent: "center",
+							backgroundColor: "var(--vscode-editor-background)",
+							borderRadius: "4px",
+							padding: "20px",
+						}}>
+						<span style={{ color: "var(--vscode-errorForeground)" }}>⚠️ {imageError}</span>
+					</div>
+				) : (
+					<img
+						src={imageUri}
+						alt={alt}
+						className="w-full h-auto rounded cursor-pointer"
+						onClick={handleOpenInEditor}
+						onError={handleImageError}
+						onLoad={handleImageLoad}
+						style={{
+							maxHeight: "400px",
+							objectFit: "contain",
+							backgroundColor: "var(--vscode-editor-background)",
+						}}
+					/>
+				)}
+				{imagePath && (
+					<div className="mt-1 text-xs text-vscode-descriptionForeground">{formatDisplayPath(imagePath)}</div>
+				)}
 				{showControls && isHovering && (
 					<div className="absolute bottom-2 right-2 flex gap-1 bg-vscode-editor-background/90 rounded p-0.5 z-10 opacity-100 transition-opacity duration-200 ease-in-out">
 						<MermaidActionButtons
@@ -202,7 +277,7 @@ export function ImageViewer({
 						onMouseUp={() => setIsDragging(false)}
 						onMouseLeave={() => setIsDragging(false)}>
 						<img
-							src={imageData}
+							src={imageUri}
 							alt={alt}
 							style={{
 								maxWidth: "90vw",
@@ -225,7 +300,7 @@ export function ImageViewer({
 						zoomInStep={0.2}
 						zoomOutStep={-0.2}
 					/>
-					{path && (
+					{imagePath && (
 						<StandardTooltip content={t("common:mermaid.buttons.copy")}>
 							<IconButton icon={copyFeedback ? "check" : "copy"} onClick={handleCopy} />
 						</StandardTooltip>