feat: optimize memory usage for image handling in webview

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
@@ -239,12 +238,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

@@ -658,9 +658,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 =
@@ -2399,6 +2407,32 @@ export class ClineProvider
 			})
 		}
 	}
+
+	/**
+	 * 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
+	 */
+	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()
+			}
+
+			// Fallback to file URI if no webview available
+			return fileUri.toString()
+		} catch (error) {
+			console.error("Failed to convert to webview URI:", error)
+			// Return file URI as fallback
+			return vscode.Uri.file(filePath).toString()
+		}
+	}
 }
 
 class OrganizationAllowListViolationError extends Error {

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,14 @@ 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 = JSON.parse(message.text || "{}")
+					return (
+						<div style={{ marginTop: "10px" }}>
+							<ImageBlock imageUri={imageInfo.imageUri} imagePath={imageInfo.imagePath} />
+						</div>
+					)
 				default:
 					return (
 						<>

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

@@ -2,14 +2,42 @@ import React from "react"
 import { ImageViewer } from "./ImageViewer"
 
 interface ImageBlockProps {
-	imageData: string
-	path?: string
+	// For new image generation tool format (preferred)
+	imageUri?: string // The webview-accessible URI for rendering
+	imagePath?: string // The actual file path for display and opening
+
+	// For backward compatibility with Mermaid diagrams and old format
+	imageData?: string // Base64 data or regular URL (legacy)
+	path?: string // Optional path for Mermaid diagrams (legacy)
 }
 
-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>
 	)
 }