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
This commit is contained in:
Roo Code 2026-01-22 11:27:05 +00:00
parent 953f037c7f
commit f3eaa464fa
6 changed files with 327 additions and 26 deletions

View file

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

View file

@ -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<void> {
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<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
}

View file

@ -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) => `<image_${index + 1}>\n${img}\n</image_${index + 1}>`)
const savedImagePaths = await this.saveImagesToTempStorage(task, images, serverName, toolName)
const imagePathsSection = savedImagePaths
.map(
(imgPath, index) =>
`<image_${index + 1}>\n <source_path>${imgPath}</source_path>\n</image_${index + 1}>`,
)
.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<string[]> {
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<string[]> {
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()

View file

@ -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("<source_path>"),
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("<source_path>"),
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"],
)

View file

@ -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."
}
}

View file

@ -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<string, unknown> }
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
}