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

packages/types/src/message.ts

@@ -146,6 +146,7 @@ export const clineSays = [
 	"api_req_retry_delayed",
 	"api_req_deleted",
 	"text",
+	"image",
 	"reasoning",
 	"completion_result",
 	"user_feedback",

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

@@ -600,9 +600,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 =
@@ -2466,4 +2474,40 @@ export class ClineProvider
 	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()
+		}
+	}
 }

src/integrations/misc/image-handler.ts

@@ -4,8 +4,46 @@ import * as vscode from "vscode"
 import { getWorkspacePath } from "../../utils/path"
 import { t } from "../../i18n"
 
-export async function openImage(dataUri: string, options?: { values?: { action?: string } }) {
-	const matches = dataUri.match(/^data:image\/([a-zA-Z]+);base64,(.+)$/)
+export async function openImage(dataUriOrPath: string, options?: { values?: { action?: string } }) {
+	// Check if it's a file path (absolute or relative)
+	const isFilePath =
+		!dataUriOrPath.startsWith("data:") &&
+		!dataUriOrPath.startsWith("http:") &&
+		!dataUriOrPath.startsWith("https:") &&
+		!dataUriOrPath.startsWith("vscode-resource:") &&
+		!dataUriOrPath.startsWith("file+.vscode-resource")
+
+	if (isFilePath) {
+		// Handle file path - open directly in VSCode
+		try {
+			// Resolve the path relative to workspace if needed
+			let filePath = dataUriOrPath
+			if (!path.isAbsolute(filePath)) {
+				const workspacePath = getWorkspacePath()
+				if (workspacePath) {
+					filePath = path.join(workspacePath, filePath)
+				}
+			}
+
+			const fileUri = vscode.Uri.file(filePath)
+
+			// Check if this is a copy action
+			if (options?.values?.action === "copy") {
+				await vscode.env.clipboard.writeText(filePath)
+				vscode.window.showInformationMessage(t("common:info.path_copied_to_clipboard"))
+				return
+			}
+
+			// Open the image file directly
+			await vscode.commands.executeCommand("vscode.open", fileUri)
+		} catch (error) {
+			vscode.window.showErrorMessage(t("common:errors.error_opening_image", { error }))
+		}
+		return
+	}
+
+	// Handle data URI (existing logic)
+	const matches = dataUriOrPath.match(/^data:image\/([a-zA-Z]+);base64,(.+)$/)
 	if (!matches) {
 		vscode.window.showErrorMessage(t("common:errors.invalid_data_uri"))
 		return

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>
 	)
 }