From f3eaa464fa0bd2e5e66ed5423998f13900c3715e Mon Sep 17 00:00:00 2001 From: Roo Code Date: Thu, 22 Jan 2026 11:27:05 +0000 Subject: [PATCH] fix: implement intermediate file persistence for MCP images - Save MCP tool images to temp storage instead of passing raw base64 to LLM - Add source_path parameter to save_image tool for copying from temp storage - This prevents data corruption and reduces token costs significantly - Images are still stored as data URLs in message.images for UI thumbnails - Update tests to match new behavior with file paths instead of base64 in text --- .../prompts/tools/native-tools/save_image.ts | 36 +++-- src/core/tools/SaveImageTool.ts | 151 +++++++++++++++++- src/core/tools/UseMcpToolTool.ts | 120 +++++++++++++- .../tools/__tests__/useMcpToolTool.spec.ts | 38 ++++- src/i18n/locales/en/tools.json | 5 +- src/shared/tools.ts | 3 +- 6 files changed, 327 insertions(+), 26 deletions(-) diff --git a/src/core/prompts/tools/native-tools/save_image.ts b/src/core/prompts/tools/native-tools/save_image.ts index 314c8ad0b2..49165e5e24 100644 --- a/src/core/prompts/tools/native-tools/save_image.ts +++ b/src/core/prompts/tools/native-tools/save_image.ts @@ -1,27 +1,36 @@ 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. +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 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. +- 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. -Example: Saving a PNG image -{ "path": "images/screenshot.png", "data": "data:image/png;base64,iVBORw0KGgoAAAANSUhEU..." } +NOTE: Either source_path OR data must be provided. -Example: Saving a JPEG image to a specific location -{ "path": "assets/captured-image", "data": "data:image/jpeg;base64,/9j/4AAQSkZJRg..." }` +Example: Saving an image from MCP tool (PREFERRED) +{ "path": "images/figma-screenshot.png", "source_path": "/path/to/temp/figma_get_screenshot_123.png" } -const PATH_PARAMETER_DESCRIPTION = `Filesystem path (relative to the workspace) where the image should be saved` +Example: Saving a base64 image (fallback) +{ "path": "images/screenshot.png", "data": "data:image/png;base64,iVBORw0KGgoAAAANSUhEU..." }` -const DATA_PARAMETER_DESCRIPTION = `Base64-encoded image data URL (e.g., 'data:image/png;base64,...')` +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: true, + strict: false, // Changed to non-strict to allow optional parameters parameters: { type: "object", properties: { @@ -29,13 +38,16 @@ export default { type: "string", description: PATH_PARAMETER_DESCRIPTION, }, + source_path: { + type: "string", + description: SOURCE_PATH_PARAMETER_DESCRIPTION, + }, data: { type: "string", description: DATA_PARAMETER_DESCRIPTION, }, }, - required: ["path", "data"], - additionalProperties: false, + required: ["path"], }, }, } satisfies OpenAI.Chat.ChatCompletionTool diff --git a/src/core/tools/SaveImageTool.ts b/src/core/tools/SaveImageTool.ts index aac282b567..fa6c8501f1 100644 --- a/src/core/tools/SaveImageTool.ts +++ b/src/core/tools/SaveImageTool.ts @@ -5,20 +5,22 @@ 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 + 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 { - const { path: relPath, data } = params + const { path: relPath, data, source_path: sourcePath } = params const { handleError, pushToolResult, askApproval } = callbacks // Validate required parameters @@ -29,15 +31,35 @@ export class SaveImageTool extends BaseTool<"save_image"> { return } - if (!data) { + // Need either source_path or data + if (!sourcePath && !data) { task.consecutiveMistakeCount++ task.recordToolError("save_image") - pushToolResult(await task.sayAndCreateMissingParamError("save_image", "data")) + 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,(.+)$/) + 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 @@ -129,6 +151,125 @@ export class SaveImageTool extends BaseTool<"save_image"> { } } + /** + * 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 { + 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 { return } diff --git a/src/core/tools/UseMcpToolTool.ts b/src/core/tools/UseMcpToolTool.ts index 6e819825af..4534a8fed8 100644 --- a/src/core/tools/UseMcpToolTool.ts +++ b/src/core/tools/UseMcpToolTool.ts @@ -1,9 +1,12 @@ +import path from "path" +import fs from "fs/promises" import type { ClineAskUseMcpServer, McpExecutionStatus } from "@roo-code/types" import { Task } from "../task/Task" import { formatResponse } from "../prompts/responses" import { t } from "../../i18n" import type { ToolUse } from "../../shared/tools" +import { getTaskDirectoryPath } from "../../utils/storage" import { BaseTool, ToolCallbacks } from "./BaseTool" @@ -322,12 +325,17 @@ export class UseMcpToolTool extends BaseTool<"use_mcp_tool"> { // Build the result text let resultText = outputText || "" - // Include image data URLs in the text response so the agent can use them with save_image tool + // If there are images, save them to temp storage and provide file paths to the LLM + // This avoids passing raw base64 through LLM context which causes corruption and high costs if (images.length > 0) { - const imageDataSection = images - .map((img, index) => `\n${img}\n`) + const savedImagePaths = await this.saveImagesToTempStorage(task, images, serverName, toolName) + const imagePathsSection = savedImagePaths + .map( + (imgPath, index) => + `\n ${imgPath}\n`, + ) .join("\n\n") - const imageInfo = `\n\n[${images.length} image(s) received - data URLs provided below for use with save_image tool]\n\n${imageDataSection}` + const imageInfo = `\n\n[${images.length} image(s) received and saved to temporary storage. Use save_image tool with source_path to save to your desired location.]\n\n${imagePathsSection}` resultText = resultText ? resultText + imageInfo : imageInfo.trim() } @@ -353,6 +361,110 @@ export class UseMcpToolTool extends BaseTool<"use_mcp_tool"> { await task.say("mcp_server_response", toolResultPretty, images) pushToolResult(formatResponse.toolResult(toolResultPretty, images)) } + + /** + * Save images to task-specific temp storage and return file paths. + * This allows passing file paths to the LLM instead of raw base64 data, + * which prevents data corruption and reduces token costs. + */ + private async saveImagesToTempStorage( + task: Task, + images: string[], + serverName: string, + toolName: string, + ): Promise { + const savedPaths: string[] = [] + + try { + const provider = task.providerRef.deref() + if (!provider) { + // Fall back to using task.cwd as temp location + return this.saveImagesToFallbackLocation(task, images, serverName, toolName) + } + + const globalStoragePath = provider.context?.globalStorageUri?.fsPath + if (!globalStoragePath) { + return this.saveImagesToFallbackLocation(task, images, serverName, toolName) + } + + // Create a temp directory for MCP images within the task directory + const taskDir = await getTaskDirectoryPath(globalStoragePath, task.taskId) + const mcpImagesDir = path.join(taskDir, "mcp_images") + await fs.mkdir(mcpImagesDir, { recursive: true }) + + const timestamp = Date.now() + + for (let i = 0; i < images.length; i++) { + const imageDataUrl = images[i] + const { format, data } = this.parseImageDataUrl(imageDataUrl) + + if (data) { + const filename = `${serverName}_${toolName}_${timestamp}_${i + 1}.${format}` + const filePath = path.join(mcpImagesDir, filename) + + const imageBuffer = Buffer.from(data, "base64") + await fs.writeFile(filePath, imageBuffer) + + savedPaths.push(filePath) + } + } + } catch (error) { + console.error("Error saving images to temp storage:", error) + // Return empty paths array on error - the LLM will see the error and handle accordingly + } + + return savedPaths + } + + /** + * Fallback method to save images to workspace .roo/temp directory + */ + private async saveImagesToFallbackLocation( + task: Task, + images: string[], + serverName: string, + toolName: string, + ): Promise { + const savedPaths: string[] = [] + + try { + const tempDir = path.join(task.cwd, ".roo", "temp", "mcp_images") + await fs.mkdir(tempDir, { recursive: true }) + + const timestamp = Date.now() + + for (let i = 0; i < images.length; i++) { + const imageDataUrl = images[i] + const { format, data } = this.parseImageDataUrl(imageDataUrl) + + if (data) { + const filename = `${serverName}_${toolName}_${timestamp}_${i + 1}.${format}` + const filePath = path.join(tempDir, filename) + + const imageBuffer = Buffer.from(data, "base64") + await fs.writeFile(filePath, imageBuffer) + + savedPaths.push(filePath) + } + } + } catch (error) { + console.error("Error saving images to fallback location:", error) + } + + return savedPaths + } + + /** + * Parse a data URL to extract format and base64 data + */ + private parseImageDataUrl(dataUrl: string): { format: string; data: string | null } { + const match = dataUrl.match(/^data:image\/(png|jpeg|jpg|gif|webp|svg\+xml);base64,(.+)$/) + if (match) { + const format = match[1] === "jpeg" ? "jpg" : match[1] === "svg+xml" ? "svg" : match[1] + return { format, data: match[2] } + } + return { format: "png", data: null } + } } export const useMcpToolTool = new UseMcpToolTool() diff --git a/src/core/tools/__tests__/useMcpToolTool.spec.ts b/src/core/tools/__tests__/useMcpToolTool.spec.ts index f44266c763..fc81b17978 100644 --- a/src/core/tools/__tests__/useMcpToolTool.spec.ts +++ b/src/core/tools/__tests__/useMcpToolTool.spec.ts @@ -4,6 +4,19 @@ import { useMcpToolTool } from "../UseMcpToolTool" import { Task } from "../../task/Task" import { ToolUse } from "../../../shared/tools" +// Mock fs/promises +vi.mock("fs/promises", () => ({ + default: { + mkdir: vi.fn().mockResolvedValue(undefined), + writeFile: vi.fn().mockResolvedValue(undefined), + }, +})) + +// Mock storage utils +vi.mock("../../../utils/storage", () => ({ + getTaskDirectoryPath: vi.fn().mockResolvedValue("/mock/storage/tasks/test-task"), +})) + // Mock dependencies vi.mock("../../prompts/responses", () => ({ formatResponse: { @@ -62,6 +75,11 @@ describe("useMcpToolTool", () => { getAllServers: vi.fn().mockReturnValue([]), }), postMessageToWebview: vi.fn(), + context: { + globalStorageUri: { + fsPath: "/mock/global/storage", + }, + }, }), } @@ -73,6 +91,8 @@ describe("useMcpToolTool", () => { ask: vi.fn(), lastMessageTs: 123456789, providerRef: mockProviderRef, + taskId: "test-task-123", + cwd: "/test/workspace", } }) @@ -638,10 +658,16 @@ describe("useMcpToolTool", () => { expect(mockTask.say).toHaveBeenCalledWith( "mcp_server_response", expect.stringContaining( - "[1 image(s) received - data URLs provided below for use with save_image tool]", + "[1 image(s) received and saved to temporary storage. Use save_image tool with source_path to save to your desired location.]", ), ["data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJ"], ) + // Text response should contain source_path XML tags, not raw base64 + expect(mockTask.say).toHaveBeenCalledWith( + "mcp_server_response", + expect.stringContaining(""), + expect.anything(), + ) expect(mockPushToolResult).toHaveBeenCalledWith(expect.stringContaining("with 1 image(s)")) }) @@ -700,6 +726,12 @@ describe("useMcpToolTool", () => { expect.stringContaining("Node name: Button"), ["data:image/png;base64,base64imagedata"], ) + // Text response should contain source_path, not raw base64 + expect(mockTask.say).toHaveBeenCalledWith( + "mcp_server_response", + expect.stringContaining(""), + expect.anything(), + ) expect(mockPushToolResult).toHaveBeenCalledWith(expect.stringContaining("with 1 image(s)")) }) @@ -756,7 +788,7 @@ describe("useMcpToolTool", () => { expect(mockTask.say).toHaveBeenCalledWith( "mcp_server_response", expect.stringContaining( - "[1 image(s) received - data URLs provided below for use with save_image tool]", + "[1 image(s) received and saved to temporary storage. Use save_image tool with source_path to save to your desired location.]", ), ["data:image/jpeg;base64,/9j/4AAQSkZJRg=="], ) @@ -819,7 +851,7 @@ describe("useMcpToolTool", () => { expect(mockTask.say).toHaveBeenCalledWith( "mcp_server_response", expect.stringContaining( - "[2 image(s) received - data URLs provided below for use with save_image tool]", + "[2 image(s) received and saved to temporary storage. Use save_image tool with source_path to save to your desired location.]", ), ["data:image/png;base64,image1data", "data:image/png;base64,image2data"], ) diff --git a/src/i18n/locales/en/tools.json b/src/i18n/locales/en/tools.json index 4f2fc57849..75ba162dd4 100644 --- a/src/i18n/locales/en/tools.json +++ b/src/i18n/locales/en/tools.json @@ -29,6 +29,9 @@ } }, "saveImage": { - "invalidDataFormat": "Invalid image data format. Expected a base64 data URL (e.g., 'data:image/png;base64,...')." + "invalidDataFormat": "Invalid image data format. Expected a base64 data URL (e.g., 'data:image/png;base64,...').", + "missingSourceOrData": "Either 'source_path' or 'data' parameter is required. Use 'source_path' for images from MCP tools, or 'data' for base64 data URLs.", + "sourceNotFound": "Source image not found at path: {{path}}", + "invalidSourceFormat": "Invalid source image format. Supported formats: PNG, JPG, JPEG, GIF, WEBP, SVG." } } diff --git a/src/shared/tools.ts b/src/shared/tools.ts index 4209d1cb60..ed5b3ef6c8 100644 --- a/src/shared/tools.ts +++ b/src/shared/tools.ts @@ -68,6 +68,7 @@ export const toolParamNames = [ "prompt", "image", "data", // save_image parameter for base64 image data + "source_path", // save_image parameter for copying from temp storage "files", // Native protocol parameter for read_file "operations", // search_and_replace parameter for multiple operations "patch", // apply_patch parameter @@ -109,7 +110,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 } + save_image: { path: string; data?: string; source_path?: string } // Add more tools as they are migrated to native protocol }