feat: add MCP image preview thumbnails and save_image tool

packages/types/src/tool.ts

@@ -37,6 +37,7 @@ export const toolNames = [
 	"update_todo_list",
 	"run_slash_command",
 	"generate_image",
+	"save_image",
 	"custom_tool",
 ] as const
 

src/core/assistant-message/presentAssistantMessage.ts

@@ -36,6 +36,7 @@ import { newTaskTool } from "../tools/NewTaskTool"
 import { updateTodoListTool } from "../tools/UpdateTodoListTool"
 import { runSlashCommandTool } from "../tools/RunSlashCommandTool"
 import { generateImageTool } from "../tools/GenerateImageTool"
+import { saveImageTool } from "../tools/SaveImageTool"
 import { applyDiffTool as applyDiffToolClass } from "../tools/ApplyDiffTool"
 import { isValidToolName, validateToolUse } from "../tools/validateToolUse"
 import { codebaseSearchTool } from "../tools/CodebaseSearchTool"
@@ -411,6 +412,8 @@ export async function presentAssistantMessage(cline: Task) {
 						return `[${block.name} for '${block.params.command}'${block.params.args ? ` with args: ${block.params.args}` : ""}]`
 					case "generate_image":
 						return `[${block.name} for '${block.params.path}']`
+					case "save_image":
+						return `[${block.name} for '${block.params.path}']`
 					default:
 						return `[${block.name}]`
 				}
@@ -919,6 +922,14 @@ export async function presentAssistantMessage(cline: Task) {
 						pushToolResult,
 					})
 					break
+				case "save_image":
+					await checkpointSaveAndMark(cline)
+					await saveImageTool.handle(cline, block as ToolUse<"save_image">, {
+						askApproval,
+						handleError,
+						pushToolResult,
+					})
+					break
 				default: {
 					// Handle unknown/invalid tool names OR custom tools
 					// This is critical for native tool calling where every tool_use MUST have a tool_result
@@ -1095,6 +1106,7 @@ function containsXmlToolMarkup(text: string): boolean {
 		"list_files",
 		"new_task",
 		"read_file",
+		"save_image",
 		"search_and_replace",
 		"search_files",
 		"search_replace",

src/core/prompts/tools/native-tools/index.ts

@@ -9,6 +9,7 @@ import codebaseSearch from "./codebase_search"
 import executeCommand from "./execute_command"
 import fetchInstructions from "./fetch_instructions"
 import generateImage from "./generate_image"
+import saveImage from "./save_image"
 import listFiles from "./list_files"
 import newTask from "./new_task"
 import { createReadFileTool, type ReadFileToolOptions } from "./read_file"
@@ -63,6 +64,7 @@ export function getNativeTools(options: NativeToolsOptions = {}): OpenAI.Chat.Ch
 		executeCommand,
 		fetchInstructions,
 		generateImage,
+		saveImage,
 		listFiles,
 		newTask,
 		createReadFileTool(readFileOptions),

src/core/prompts/tools/native-tools/save_image.ts

@@ -0,0 +1,53 @@
+import type OpenAI from "openai"
+
+const SAVE_IMAGE_DESCRIPTION = `Request to save an image to a file. This tool supports two methods:
+
+1. **Using source_path (PREFERRED for MCP tools)**: When you receive images from MCP tools like Figma, the images are automatically saved to temporary storage and you receive file paths. Use the source_path parameter to copy the image to your desired location. This is efficient and avoids data corruption.
+
+2. **Using data (for base64 data URLs)**: For images provided as base64 data URLs from other sources.
+
+Parameters:
+- path: (required) The destination file path where the image should be saved (relative to the current workspace directory). The tool will automatically add the appropriate image extension based on the source image format if not provided.
+- source_path: (optional) The absolute path to a source image file (typically from MCP tool temporary storage). Use this for images received from MCP tools - the path is provided in the tool response. PREFERRED over data.
+- data: (optional) Base64-encoded image data URL (e.g., 'data:image/png;base64,...'). Supported formats: PNG, JPG, JPEG, GIF, WEBP, SVG. Only use if source_path is not available.
+
+NOTE: Either source_path OR data must be provided.
+
+Example: Saving an image from MCP tool (PREFERRED)
+{ "path": "images/figma-screenshot.png", "source_path": "/path/to/temp/figma_get_screenshot_123.png" }
+
+Example: Saving a base64 image (fallback)
+{ "path": "images/screenshot.png", "data": "data:image/png;base64,iVBORw0KGgoAAAANSUhEU..." }`
+
+const PATH_PARAMETER_DESCRIPTION = `Destination filesystem path (relative to the workspace) where the image should be saved`
+
+const SOURCE_PATH_PARAMETER_DESCRIPTION = `Absolute path to a source image file (from MCP tool temporary storage). PREFERRED method for saving images from MCP tools.`
+
+const DATA_PARAMETER_DESCRIPTION = `Base64-encoded image data URL (e.g., 'data:image/png;base64,...'). Only use if source_path is not available.`
+
+export default {
+	type: "function",
+	function: {
+		name: "save_image",
+		description: SAVE_IMAGE_DESCRIPTION,
+		strict: false, // Changed to non-strict to allow optional parameters
+		parameters: {
+			type: "object",
+			properties: {
+				path: {
+					type: "string",
+					description: PATH_PARAMETER_DESCRIPTION,
+				},
+				source_path: {
+					type: "string",
+					description: SOURCE_PATH_PARAMETER_DESCRIPTION,
+				},
+				data: {
+					type: "string",
+					description: DATA_PARAMETER_DESCRIPTION,
+				},
+			},
+			required: ["path"],
+		},
+	},
+} satisfies OpenAI.Chat.ChatCompletionTool

src/core/tools/SaveImageTool.ts

@@ -0,0 +1,278 @@
+import path from "path"
+import fs from "fs/promises"
+import * as vscode from "vscode"
+import { Task } from "../task/Task"
+import { formatResponse } from "../prompts/responses"
+import { getReadablePath } from "../../utils/path"
+import { isPathOutsideWorkspace } from "../../utils/pathUtils"
+import { fileExistsAtPath } from "../../utils/fs"
+import { BaseTool, ToolCallbacks } from "./BaseTool"
+import type { ToolUse } from "../../shared/tools"
+import { t } from "../../i18n"
+
+interface SaveImageParams {
+	path: string
+	data?: string
+	source_path?: string
+}
+
+export class SaveImageTool extends BaseTool<"save_image"> {
+	readonly name = "save_image" as const
+
+	async execute(params: SaveImageParams, task: Task, callbacks: ToolCallbacks): Promise<void> {
+		const { path: relPath, data, source_path: sourcePath } = params
+		const { handleError, pushToolResult, askApproval } = callbacks
+
+		// Validate required parameters
+		if (!relPath) {
+			task.consecutiveMistakeCount++
+			task.recordToolError("save_image")
+			pushToolResult(await task.sayAndCreateMissingParamError("save_image", "path"))
+			return
+		}
+
+		// Need either source_path or data
+		if (!sourcePath && !data) {
+			task.consecutiveMistakeCount++
+			task.recordToolError("save_image")
+			await task.say(
+				"error",
+				t("tools:saveImage.missingSourceOrData", {
+					defaultValue:
+						"Either 'source_path' or 'data' parameter is required. Use 'source_path' for images from MCP tools, or 'data' for base64 data URLs.",
+				}),
+			)
+			task.didToolFailInCurrentTurn = true
+			pushToolResult(
+				formatResponse.toolError(
+					"Either 'source_path' or 'data' parameter is required. Use 'source_path' for images from MCP tools, or 'data' for base64 data URLs.",
+				),
+			)
+			return
+		}
+
+		// If source_path is provided, use it to copy the file
+		if (sourcePath) {
+			await this.copyFromSourcePath(task, sourcePath, relPath, callbacks)
+			return
+		}
+
+		// Otherwise, use the data parameter (base64 data URL)
+		// Validate the image data format first (to determine finalPath)
+		const base64Match = data!.match(/^data:image\/(png|jpeg|jpg|gif|webp|svg\+xml);base64,(.+)$/)
+		if (!base64Match) {
+			await task.say("error", t("tools:saveImage.invalidDataFormat"))
+			task.didToolFailInCurrentTurn = true
+			pushToolResult(
+				formatResponse.toolError(
+					"Invalid image data format. Expected a base64 data URL (e.g., 'data:image/png;base64,...').",
+				),
+			)
+			return
+		}
+
+		const imageFormat = base64Match[1]
+		const base64Data = base64Match[2]
+
+		// Ensure the path has a valid image extension
+		let finalPath = relPath
+		if (!finalPath.match(/\.(png|jpg|jpeg|gif|webp|svg)$/i)) {
+			// Add extension based on the data format
+			const ext = imageFormat === "jpeg" ? "jpg" : imageFormat === "svg+xml" ? "svg" : imageFormat
+			finalPath = `${finalPath}.${ext}`
+		}
+
+		// Validate access via .rooignore (using finalPath after extension is added)
+		const accessAllowed = task.rooIgnoreController?.validateAccess(finalPath)
+		if (!accessAllowed) {
+			await task.say("rooignore_error", finalPath)
+			pushToolResult(formatResponse.rooIgnoreError(finalPath))
+			return
+		}
+
+		// Check write protection (using finalPath after extension is added)
+		const isWriteProtected = task.rooProtectedController?.isWriteProtected(finalPath) || false
+
+		const fullPath = path.resolve(task.cwd, finalPath)
+		const isOutsideWorkspace = isPathOutsideWorkspace(fullPath)
+
+		const sharedMessageProps = {
+			tool: "saveImage" as const,
+			path: getReadablePath(task.cwd, finalPath),
+			isOutsideWorkspace,
+			isProtected: isWriteProtected,
+		}
+
+		try {
+			task.consecutiveMistakeCount = 0
+
+			const approvalMessage = JSON.stringify({
+				...sharedMessageProps,
+				content: `Save image to ${getReadablePath(task.cwd, finalPath)}`,
+			})
+
+			const didApprove = await askApproval("tool", approvalMessage, undefined, isWriteProtected)
+
+			if (!didApprove) {
+				return
+			}
+
+			// Convert base64 to buffer and save
+			const imageBuffer = Buffer.from(base64Data, "base64")
+
+			const absolutePath = path.resolve(task.cwd, finalPath)
+			const directory = path.dirname(absolutePath)
+			await fs.mkdir(directory, { recursive: true })
+
+			await fs.writeFile(absolutePath, imageBuffer)
+
+			// Track the file context
+			if (finalPath) {
+				await task.fileContextTracker.trackFileContext(finalPath, "roo_edited")
+			}
+
+			task.didEditFile = true
+
+			task.recordToolUsage("save_image")
+
+			const provider = task.providerRef.deref()
+			const fullImagePath = path.join(task.cwd, finalPath)
+
+			let imageUri = provider?.convertToWebviewUri?.(fullImagePath) ?? vscode.Uri.file(fullImagePath).toString()
+
+			// Add cache buster to force refresh
+			const cacheBuster = Date.now()
+			imageUri = imageUri.includes("?") ? `${imageUri}&t=${cacheBuster}` : `${imageUri}?t=${cacheBuster}`
+
+			await task.say("image", JSON.stringify({ imageUri, imagePath: fullImagePath }))
+			pushToolResult(formatResponse.toolResult(`Image saved to ${getReadablePath(task.cwd, finalPath)}`))
+		} catch (error) {
+			await handleError("saving image", error as Error)
+		}
+	}
+
+	/**
+	 * Copy an image from a source path (typically from MCP temp storage) to the destination path.
+	 * This is the preferred method for saving images from MCP tools as it avoids passing
+	 * raw base64 through LLM context.
+	 */
+	private async copyFromSourcePath(
+		task: Task,
+		sourcePath: string,
+		destRelPath: string,
+		callbacks: ToolCallbacks,
+	): Promise<void> {
+		const { handleError, pushToolResult, askApproval } = callbacks
+
+		try {
+			// Check if source file exists
+			const sourceExists = await fileExistsAtPath(sourcePath)
+			if (!sourceExists) {
+				task.consecutiveMistakeCount++
+				task.recordToolError("save_image")
+				await task.say(
+					"error",
+					t("tools:saveImage.sourceNotFound", {
+						defaultValue: `Source image not found at path: ${sourcePath}`,
+						path: sourcePath,
+					}),
+				)
+				task.didToolFailInCurrentTurn = true
+				pushToolResult(formatResponse.toolError(`Source image not found at path: ${sourcePath}`))
+				return
+			}
+
+			// Get extension from source file
+			const sourceExt = path.extname(sourcePath).toLowerCase()
+			const validExtensions = [".png", ".jpg", ".jpeg", ".gif", ".webp", ".svg"]
+
+			if (!validExtensions.includes(sourceExt)) {
+				task.consecutiveMistakeCount++
+				task.recordToolError("save_image")
+				await task.say("error", t("tools:saveImage.invalidSourceFormat"))
+				task.didToolFailInCurrentTurn = true
+				pushToolResult(
+					formatResponse.toolError(
+						`Invalid source image format. Supported formats: ${validExtensions.join(", ")}`,
+					),
+				)
+				return
+			}
+
+			// Ensure the destination path has the correct extension
+			let finalPath = destRelPath
+			if (!finalPath.match(/\.(png|jpg|jpeg|gif|webp|svg)$/i)) {
+				finalPath = `${finalPath}${sourceExt}`
+			}
+
+			// Validate access via .rooignore
+			const accessAllowed = task.rooIgnoreController?.validateAccess(finalPath)
+			if (!accessAllowed) {
+				await task.say("rooignore_error", finalPath)
+				pushToolResult(formatResponse.rooIgnoreError(finalPath))
+				return
+			}
+
+			// Check write protection
+			const isWriteProtected = task.rooProtectedController?.isWriteProtected(finalPath) || false
+
+			const fullPath = path.resolve(task.cwd, finalPath)
+			const isOutsideWorkspace = isPathOutsideWorkspace(fullPath)
+
+			const sharedMessageProps = {
+				tool: "saveImage" as const,
+				path: getReadablePath(task.cwd, finalPath),
+				isOutsideWorkspace,
+				isProtected: isWriteProtected,
+			}
+
+			task.consecutiveMistakeCount = 0
+
+			const approvalMessage = JSON.stringify({
+				...sharedMessageProps,
+				content: `Save image from ${sourcePath} to ${getReadablePath(task.cwd, finalPath)}`,
+			})
+
+			const didApprove = await askApproval("tool", approvalMessage, undefined, isWriteProtected)
+
+			if (!didApprove) {
+				return
+			}
+
+			// Create destination directory and copy file
+			const absolutePath = path.resolve(task.cwd, finalPath)
+			const directory = path.dirname(absolutePath)
+			await fs.mkdir(directory, { recursive: true })
+
+			await fs.copyFile(sourcePath, absolutePath)
+
+			// Track the file context
+			if (finalPath) {
+				await task.fileContextTracker.trackFileContext(finalPath, "roo_edited")
+			}
+
+			task.didEditFile = true
+			task.recordToolUsage("save_image")
+
+			const provider = task.providerRef.deref()
+			const fullImagePath = path.join(task.cwd, finalPath)
+
+			let imageUri = provider?.convertToWebviewUri?.(fullImagePath) ?? vscode.Uri.file(fullImagePath).toString()
+
+			// Add cache buster to force refresh
+			const cacheBuster = Date.now()
+			imageUri = imageUri.includes("?") ? `${imageUri}&t=${cacheBuster}` : `${imageUri}?t=${cacheBuster}`
+
+			await task.say("image", JSON.stringify({ imageUri, imagePath: fullImagePath }))
+			pushToolResult(formatResponse.toolResult(`Image saved to ${getReadablePath(task.cwd, finalPath)}`))
+		} catch (error) {
+			await handleError("saving image", error as Error)
+		}
+	}
+
+	override async handlePartial(task: Task, block: ToolUse<"save_image">): Promise<void> {
+		return
+	}
+}
+
+export const saveImageTool = new SaveImageTool()