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
This commit is contained in:
Daniel 2025-08-30 14:57:30 -05:00 committed by Shawn
parent 08d3b6c5ab
commit 68b68a22cd
26 changed files with 417 additions and 58 deletions

View file

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

View file

@ -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

View file

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

View file

@ -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

View file

@ -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 (
<>

View file

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

View file

@ -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>

View file

@ -0,0 +1,116 @@
// npx vitest run src/components/common/__tests__/ImageViewer.spec.tsx
import { render, fireEvent } from "@testing-library/react"
import { describe, it, expect, vi } from "vitest"
import { ImageViewer } from "../ImageViewer"
// Mock vscode API
vi.mock("@src/utils/vscode", () => ({
vscode: {
postMessage: vi.fn(),
},
}))
// Import the mocked vscode after the mock is set up
import { vscode } from "@src/utils/vscode"
describe("ImageViewer", () => {
it("should render image with webview URI", () => {
const webviewUri = "https://file+.vscode-resource.vscode-cdn.net/path/to/image.png"
const { container } = render(<ImageViewer imageUri={webviewUri} alt="Test image" />)
const img = container.querySelector("img")
expect(img).toBeTruthy()
expect(img?.src).toBe(webviewUri)
expect(img?.alt).toBe("Test image")
})
it("should render image with vscode-resource URI", () => {
const vscodeResourceUri = "vscode-resource://file///path/to/image.png"
const { container } = render(<ImageViewer imageUri={vscodeResourceUri} alt="Test image" />)
const img = container.querySelector("img")
expect(img).toBeTruthy()
expect(img?.src).toBe(vscodeResourceUri)
})
it("should handle base64 images", () => {
const base64Image =
"data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg=="
const { container } = render(<ImageViewer imageUri={base64Image} alt="Base64 image" />)
const img = container.querySelector("img")
expect(img).toBeTruthy()
expect(img?.src).toBe(base64Image)
})
it("should use imageUri for rendering and imagePath for display", () => {
const webviewUri = "https://file+.vscode-resource.vscode-cdn.net/path/to/image.png"
const filePath = "/Users/test/project/image.png"
const { container } = render(<ImageViewer imageUri={webviewUri} imagePath={filePath} alt="Test image" />)
const img = container.querySelector("img")
expect(img).toBeTruthy()
// Should use imageUri for src
expect(img?.src).toBe(webviewUri)
// Should display imagePath below image
const pathElement = container.querySelector(".text-xs.text-vscode-descriptionForeground")
expect(pathElement).toBeTruthy()
expect(pathElement?.textContent).toContain("image.png")
})
it("should handle click to open in editor", () => {
const webviewUri = "https://file+.vscode-resource.vscode-cdn.net/path/to/image.png"
const filePath = "/Users/test/project/image.png"
const { container } = render(<ImageViewer imageUri={webviewUri} imagePath={filePath} alt="Test image" />)
const img = container.querySelector("img")
expect(img).toBeTruthy()
// Clear previous calls
vi.clearAllMocks()
// Click the image
fireEvent.click(img!)
// Check if vscode.postMessage was called to open the image with the actual path
expect(vscode.postMessage).toHaveBeenCalledWith({
type: "openImage",
text: filePath,
})
})
it("should handle error state gracefully", () => {
const invalidUri = "invalid://uri"
const { container } = render(<ImageViewer imageUri={invalidUri} alt="Invalid image" />)
const img = container.querySelector("img")
expect(img).toBeTruthy()
// Trigger error event
fireEvent.error(img!)
// Image should still be rendered but might have error styling
expect(img).toBeTruthy()
})
it("should show no image message when imageUri is empty", () => {
const { container } = render(<ImageViewer imageUri="" alt="Empty image" />)
// Should show no image message
expect(container.textContent).toContain("common:image.noData")
})
it("should display path below image when provided", () => {
const filePath = "/Users/test/rc1/path/to/image.png"
const webviewUri = "https://file+.vscode-resource.vscode-cdn.net/path/to/image.png"
const { container } = render(<ImageViewer imageUri={webviewUri} imagePath={filePath} alt="Test image" />)
// Check if path is displayed as relative path
const pathElement = container.querySelector(".text-xs.text-vscode-descriptionForeground")
expect(pathElement).toBeTruthy()
// Accept filename or relative path depending on environment
expect(pathElement?.textContent).toContain("image.png")
})
})

View file

@ -51,7 +51,8 @@
"image": {
"tabs": {
"view": "Imatge"
}
},
"noData": "Sense dades d'imatge"
},
"file": {
"errors": {

View file

@ -51,7 +51,8 @@
"image": {
"tabs": {
"view": "Bild"
}
},
"noData": "Keine Bilddaten"
},
"file": {
"errors": {

View file

@ -51,7 +51,8 @@
"image": {
"tabs": {
"view": "Image"
}
},
"noData": "No image data"
},
"file": {
"errors": {

View file

@ -51,7 +51,8 @@
"image": {
"tabs": {
"view": "Imagen"
}
},
"noData": "Sin datos de imagen"
},
"file": {
"errors": {

View file

@ -51,7 +51,8 @@
"image": {
"tabs": {
"view": "Image"
}
},
"noData": "Aucune donnée d'image"
},
"file": {
"errors": {

View file

@ -51,7 +51,8 @@
"image": {
"tabs": {
"view": "चित्र"
}
},
"noData": "कोई छवि डेटा नहीं"
},
"file": {
"errors": {

View file

@ -51,7 +51,8 @@
"image": {
"tabs": {
"view": "Gambar"
}
},
"noData": "Tidak ada data gambar"
},
"file": {
"errors": {

View file

@ -51,7 +51,8 @@
"image": {
"tabs": {
"view": "Immagine"
}
},
"noData": "Nessun dato immagine"
},
"file": {
"errors": {

View file

@ -51,7 +51,8 @@
"image": {
"tabs": {
"view": "画像"
}
},
"noData": "画像データなし"
},
"file": {
"errors": {

View file

@ -51,7 +51,8 @@
"image": {
"tabs": {
"view": "이미지"
}
},
"noData": "이미지 데이터 없음"
},
"file": {
"errors": {

View file

@ -51,7 +51,8 @@
"image": {
"tabs": {
"view": "Afbeelding"
}
},
"noData": "Geen afbeeldingsgegevens"
},
"file": {
"errors": {

View file

@ -51,7 +51,8 @@
"image": {
"tabs": {
"view": "Obraz"
}
},
"noData": "Brak danych obrazu"
},
"file": {
"errors": {

View file

@ -51,7 +51,8 @@
"image": {
"tabs": {
"view": "Imagem"
}
},
"noData": "Nenhum dado de imagem"
},
"file": {
"errors": {

View file

@ -51,7 +51,8 @@
"image": {
"tabs": {
"view": "Изображение"
}
},
"noData": "Нет данных изображения"
},
"file": {
"errors": {

View file

@ -51,7 +51,8 @@
"image": {
"tabs": {
"view": "Resim"
}
},
"noData": "Resim verisi yok"
},
"file": {
"errors": {

View file

@ -51,7 +51,8 @@
"image": {
"tabs": {
"view": "Hình ảnh"
}
},
"noData": "Không có dữ liệu hình ảnh"
},
"file": {
"errors": {

View file

@ -51,7 +51,8 @@
"image": {
"tabs": {
"view": "图像"
}
},
"noData": "无图片数据"
},
"file": {
"errors": {

View file

@ -51,7 +51,8 @@
"image": {
"tabs": {
"view": "圖像"
}
},
"noData": "無圖片資料"
},
"file": {
"errors": {