From 1632659c029e956b9a0f36ca4fd0df1cd2ea6736 Mon Sep 17 00:00:00 2001 From: Roo Code Date: Wed, 21 Jan 2026 23:02:54 +0000 Subject: [PATCH] feat: add MCP image preview thumbnails and save_image tool - Add image thumbnails to McpExecution component (Feature 1) - Import Thumbnails component - Add images prop to McpExecutionProps interface - Render thumbnails when images are present (click to open in VSCode) - Pass message.images from ChatRow to McpExecution - Add save_image tool for agent to save images (Feature 2) - Create SaveImageTool.ts with base64 data URL support - Add save_image to toolNames in types/tool.ts - Add tool definitions in shared/tools.ts - Create native tool description in prompts/tools/native-tools - Register tool in presentAssistantMessage.ts - Add to edit tool group for file write operations Addresses Issue #10877 --- packages/types/src/tool.ts | 1 + .../presentAssistantMessage.ts | 12 ++ src/core/prompts/tools/native-tools/index.ts | 2 + .../prompts/tools/native-tools/save_image.ts | 41 ++++++ src/core/tools/SaveImageTool.ts | 137 ++++++++++++++++++ src/shared/tools.ts | 5 +- webview-ui/src/components/chat/ChatRow.tsx | 1 + .../src/components/chat/McpExecution.tsx | 10 ++ 8 files changed, 208 insertions(+), 1 deletion(-) create mode 100644 src/core/prompts/tools/native-tools/save_image.ts create mode 100644 src/core/tools/SaveImageTool.ts diff --git a/packages/types/src/tool.ts b/packages/types/src/tool.ts index 147eb24b6c..ed4dd60f85 100644 --- a/packages/types/src/tool.ts +++ b/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 diff --git a/src/core/assistant-message/presentAssistantMessage.ts b/src/core/assistant-message/presentAssistantMessage.ts index 6469ba8a5c..95d6771276 100644 --- a/src/core/assistant-message/presentAssistantMessage.ts +++ b/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", diff --git a/src/core/prompts/tools/native-tools/index.ts b/src/core/prompts/tools/native-tools/index.ts index 4f78729cdc..d171d83926 100644 --- a/src/core/prompts/tools/native-tools/index.ts +++ b/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), diff --git a/src/core/prompts/tools/native-tools/save_image.ts b/src/core/prompts/tools/native-tools/save_image.ts new file mode 100644 index 0000000000..314c8ad0b2 --- /dev/null +++ b/src/core/prompts/tools/native-tools/save_image.ts @@ -0,0 +1,41 @@ +import type OpenAI from "openai" + +const SAVE_IMAGE_DESCRIPTION = `Request to save a base64-encoded image to a file. This tool is useful for saving images that were received from MCP tools or other sources. The image data must be provided as a base64 data URL. + +Parameters: +- path: (required) The 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 image format if not provided. +- data: (required) The base64-encoded image data URL (e.g., 'data:image/png;base64,...'). Supported formats: PNG, JPG, JPEG, GIF, WEBP, SVG. + +Example: Saving a PNG image +{ "path": "images/screenshot.png", "data": "data:image/png;base64,iVBORw0KGgoAAAANSUhEU..." } + +Example: Saving a JPEG image to a specific location +{ "path": "assets/captured-image", "data": "data:image/jpeg;base64,/9j/4AAQSkZJRg..." }` + +const PATH_PARAMETER_DESCRIPTION = `Filesystem path (relative to the workspace) where the image should be saved` + +const DATA_PARAMETER_DESCRIPTION = `Base64-encoded image data URL (e.g., 'data:image/png;base64,...')` + +export default { + type: "function", + function: { + name: "save_image", + description: SAVE_IMAGE_DESCRIPTION, + strict: true, + parameters: { + type: "object", + properties: { + path: { + type: "string", + description: PATH_PARAMETER_DESCRIPTION, + }, + data: { + type: "string", + description: DATA_PARAMETER_DESCRIPTION, + }, + }, + required: ["path", "data"], + additionalProperties: false, + }, + }, +} satisfies OpenAI.Chat.ChatCompletionTool diff --git a/src/core/tools/SaveImageTool.ts b/src/core/tools/SaveImageTool.ts new file mode 100644 index 0000000000..ecf9fbf5b9 --- /dev/null +++ b/src/core/tools/SaveImageTool.ts @@ -0,0 +1,137 @@ +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 { BaseTool, ToolCallbacks } from "./BaseTool" +import type { ToolUse } from "../../shared/tools" +import { t } from "../../i18n" + +interface SaveImageParams { + path: string + data: string +} + +export class SaveImageTool extends BaseTool<"save_image"> { + readonly name = "save_image" as const + + async execute(params: SaveImageParams, task: Task, callbacks: ToolCallbacks): Promise { + const { path: relPath, data } = 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 + } + + if (!data) { + task.consecutiveMistakeCount++ + task.recordToolError("save_image") + pushToolResult(await task.sayAndCreateMissingParamError("save_image", "data")) + return + } + + // Validate access via .rooignore + const accessAllowed = task.rooIgnoreController?.validateAccess(relPath) + if (!accessAllowed) { + await task.say("rooignore_error", relPath) + pushToolResult(formatResponse.rooIgnoreError(relPath)) + return + } + + // Check write protection + const isWriteProtected = task.rooProtectedController?.isWriteProtected(relPath) || false + + const fullPath = path.resolve(task.cwd, relPath) + const isOutsideWorkspace = isPathOutsideWorkspace(fullPath) + + // Validate the image data format + 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}` + } + + 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) + } + } + + override async handlePartial(task: Task, block: ToolUse<"save_image">): Promise { + return + } +} + +export const saveImageTool = new SaveImageTool() diff --git a/src/shared/tools.ts b/src/shared/tools.ts index 01632b2746..4209d1cb60 100644 --- a/src/shared/tools.ts +++ b/src/shared/tools.ts @@ -67,6 +67,7 @@ export const toolParamNames = [ "todos", "prompt", "image", + "data", // save_image parameter for base64 image data "files", // Native protocol parameter for read_file "operations", // search_and_replace parameter for multiple operations "patch", // apply_patch parameter @@ -108,6 +109,7 @@ export type NativeToolArgs = { update_todo_list: { todos: string } use_mcp_tool: { server_name: string; tool_name: string; arguments?: Record } write_to_file: { path: string; content: string } + save_image: { path: string; data: string } // Add more tools as they are migrated to native protocol } @@ -264,6 +266,7 @@ export const TOOL_DISPLAY_NAMES: Record = { update_todo_list: "update todo list", run_slash_command: "run slash command", generate_image: "generate images", + save_image: "save images", custom_tool: "use custom tools", } as const @@ -273,7 +276,7 @@ export const TOOL_GROUPS: Record = { tools: ["read_file", "fetch_instructions", "search_files", "list_files", "codebase_search"], }, edit: { - tools: ["apply_diff", "write_to_file", "generate_image"], + tools: ["apply_diff", "write_to_file", "generate_image", "save_image"], customTools: ["search_and_replace", "search_replace", "edit_file", "apply_patch"], }, browser: { diff --git a/webview-ui/src/components/chat/ChatRow.tsx b/webview-ui/src/components/chat/ChatRow.tsx index e71f92dc41..c15d760cb5 100644 --- a/webview-ui/src/components/chat/ChatRow.tsx +++ b/webview-ui/src/components/chat/ChatRow.tsx @@ -1627,6 +1627,7 @@ export const ChatRowContent = ({ server={server} useMcpServer={useMcpServer} alwaysAllowMcp={alwaysAllowMcp} + images={message.images} /> )} diff --git a/webview-ui/src/components/chat/McpExecution.tsx b/webview-ui/src/components/chat/McpExecution.tsx index 9e48552fdc..c9b62ec351 100644 --- a/webview-ui/src/components/chat/McpExecution.tsx +++ b/webview-ui/src/components/chat/McpExecution.tsx @@ -16,6 +16,7 @@ import { cn } from "@src/lib/utils" import { Button } from "@src/components/ui" import CodeBlock from "../common/CodeBlock" +import Thumbnails from "../common/Thumbnails" import McpToolRow from "../mcp/McpToolRow" import { Markdown } from "./Markdown" @@ -36,6 +37,7 @@ interface McpExecutionProps { } useMcpServer?: ClineAskUseMcpServer alwaysAllowMcp?: boolean + images?: string[] } export const McpExecution = ({ @@ -47,6 +49,7 @@ export const McpExecution = ({ server, useMcpServer, alwaysAllowMcp = false, + images, }: McpExecutionProps) => { const { t } = useTranslation("mcp") @@ -289,6 +292,13 @@ export const McpExecution = ({ hasArguments={!!(isArguments || useMcpServer?.arguments || argumentsText)} isPartial={status ? status.status !== "completed" : false} /> + + {/* Images section - show thumbnails of returned images */} + {images && images.length > 0 && ( +
+ +
+ )} )