diff --git a/src/core/tools/UseMcpToolTool.ts b/src/core/tools/UseMcpToolTool.ts index b44b96054e..357de333db 100644 --- a/src/core/tools/UseMcpToolTool.ts +++ b/src/core/tools/UseMcpToolTool.ts @@ -5,6 +5,9 @@ import { McpExecutionStatus } from "@roo-code/types" import { t } from "../../i18n" import { BaseTool, ToolCallbacks } from "./BaseTool" import type { ToolUse } from "../../shared/tools" +import * as fs from "fs/promises" +import * as path from "path" +import { readImageAsDataUrlWithBuffer, isSupportedImageFormat } from "./helpers/imageHelpers" interface UseMcpToolParams { server_name: string @@ -34,6 +37,73 @@ export class UseMcpToolTool extends BaseTool<"use_mcp_tool"> { } } + /** + * Checks if a value appears to be a file path pointing to an image + */ + private async isImageFilePath(value: unknown): Promise { + if (typeof value !== "string") { + return false + } + + // Skip if it's already a base64 data URL + if (value.startsWith("data:image/")) { + return false + } + + // Check if it looks like a file path + const ext = path.extname(value).toLowerCase() + if (!isSupportedImageFormat(ext)) { + return false + } + + // Check if the file exists + try { + const stats = await fs.stat(value) + return stats.isFile() + } catch { + // File doesn't exist or can't be accessed + return false + } + } + + /** + * Recursively processes an object/array to convert image file paths to base64 data URLs + */ + private async convertImagePathsToBase64(obj: unknown): Promise { + if (obj === null || obj === undefined) { + return obj + } + + // Handle arrays + if (Array.isArray(obj)) { + return Promise.all(obj.map((item) => this.convertImagePathsToBase64(item))) + } + + // Handle objects + if (typeof obj === "object") { + const result: Record = {} + for (const [key, value] of Object.entries(obj)) { + result[key] = await this.convertImagePathsToBase64(value) + } + return result + } + + // Handle potential image file paths + if (await this.isImageFilePath(obj)) { + try { + const { dataUrl } = await readImageAsDataUrlWithBuffer(obj as string) + return dataUrl + } catch (error) { + // If we can't read the image, return the original value + console.error(`Failed to convert image file to base64: ${error}`) + return obj + } + } + + // Return primitive values as-is + return obj + } + async execute(params: UseMcpToolParams, task: Task, callbacks: ToolCallbacks): Promise { const { askApproval, handleError, pushToolResult, toolProtocol } = callbacks @@ -55,12 +125,17 @@ export class UseMcpToolTool extends BaseTool<"use_mcp_tool"> { // Reset mistake count on successful validation task.consecutiveMistakeCount = 0 + // Convert any image file paths in arguments to base64 data URLs + const processedArguments = parsedArguments + ? ((await this.convertImagePathsToBase64(parsedArguments)) as Record) + : parsedArguments + // Get user approval const completeMessage = JSON.stringify({ type: "use_mcp_tool", serverName, toolName, - arguments: params.arguments ? JSON.stringify(params.arguments) : undefined, + arguments: processedArguments ? JSON.stringify(processedArguments) : undefined, } satisfies ClineAskUseMcpServer) const executionId = task.lastMessageTs?.toString() ?? Date.now().toString() @@ -75,7 +150,7 @@ export class UseMcpToolTool extends BaseTool<"use_mcp_tool"> { task, serverName, toolName, - parsedArguments, + processedArguments, executionId, pushToolResult, ) diff --git a/src/core/tools/__tests__/useMcpToolTool.spec.ts b/src/core/tools/__tests__/useMcpToolTool.spec.ts index 130047ae15..a79091e37d 100644 --- a/src/core/tools/__tests__/useMcpToolTool.spec.ts +++ b/src/core/tools/__tests__/useMcpToolTool.spec.ts @@ -3,6 +3,25 @@ import { useMcpToolTool } from "../UseMcpToolTool" import { Task } from "../../task/Task" import { ToolUse } from "../../../shared/tools" +import * as fs from "fs/promises" + +// Mock fs/promises +vi.mock("fs/promises", () => ({ + stat: vi.fn(), + readFile: vi.fn(), +})) + +// Import the actual module to get the mock +import { readImageAsDataUrlWithBuffer } from "../helpers/imageHelpers" + +// Mock image helpers +vi.mock("../helpers/imageHelpers", () => ({ + readImageAsDataUrlWithBuffer: vi.fn(), + isSupportedImageFormat: vi.fn((ext: string) => { + const supportedFormats = [".png", ".jpg", ".jpeg", ".gif", ".webp"] + return supportedFormats.includes(ext.toLowerCase()) + }), +})) // Mock dependencies vi.mock("../../prompts/responses", () => ({ @@ -537,4 +556,347 @@ describe("useMcpToolTool", () => { expect(mockAskApproval).not.toHaveBeenCalled() }) }) + + describe("image handling", () => { + let mockReadImageAsDataUrlWithBuffer: ReturnType + let mockFsStat: ReturnType + + beforeEach(() => { + // Get mocked functions + mockReadImageAsDataUrlWithBuffer = vi.mocked(readImageAsDataUrlWithBuffer) + mockFsStat = vi.mocked(fs.stat) + + // Clear all mocks before each test + vi.clearAllMocks() + }) + + it("should convert image file paths to base64 data URLs", async () => { + // Setup + const imagePath = "/path/to/image.png" + const base64DataUrl = + "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg==" + + // Mock file exists and is a file + mockFsStat.mockResolvedValue({ + isFile: () => true, + }) + + // Mock image reading + mockReadImageAsDataUrlWithBuffer.mockResolvedValue({ + dataUrl: base64DataUrl, + buffer: Buffer.from("test"), + }) + + // Mock server and tool exist + mockProviderRef.deref.mockReturnValue({ + getMcpHub: () => ({ + getAllServers: vi + .fn() + .mockReturnValue([ + { + name: "test_server", + tools: [{ name: "process_image", description: "Process an image" }], + }, + ]), + callTool: vi.fn().mockResolvedValue({ + content: [{ type: "text", text: "Image processed" }], + isError: false, + }), + }), + postMessageToWebview: vi.fn(), + }) + + const block: ToolUse = { + type: "tool_use", + name: "use_mcp_tool", + params: { + server_name: "test_server", + tool_name: "process_image", + arguments: JSON.stringify({ image: imagePath }), + }, + partial: false, + } + + mockAskApproval.mockResolvedValue(true) + + // Execute + await useMcpToolTool.handle(mockTask as Task, block as any, { + askApproval: mockAskApproval, + handleError: mockHandleError, + pushToolResult: mockPushToolResult, + removeClosingTag: mockRemoveClosingTag, + toolProtocol: "xml", + }) + + // Verify image was converted + expect(mockFsStat).toHaveBeenCalledWith(imagePath) + expect(mockReadImageAsDataUrlWithBuffer).toHaveBeenCalledWith(imagePath) + + // Verify the approval message contains the base64 data URL + const approvalCall = mockAskApproval.mock.calls[0] + const approvalMessage = JSON.parse(approvalCall[1]) + const args = JSON.parse(approvalMessage.arguments) + expect(args.image).toBe(base64DataUrl) + }) + + it("should handle nested image paths in complex objects", async () => { + // Setup + const imagePath1 = "/path/to/image1.jpg" + const imagePath2 = "/path/to/image2.png" + const base64DataUrl1 = "data:image/jpeg;base64,/9j/4AAQSkZJRg==" + const base64DataUrl2 = "data:image/png;base64,iVBORw0KGgoAAAANS==" + + // Mock file exists checks + mockFsStat.mockImplementation((path) => { + if (path === imagePath1 || path === imagePath2) { + return Promise.resolve({ isFile: () => true }) + } + return Promise.reject(new Error("File not found")) + }) + + // Mock image reading + mockReadImageAsDataUrlWithBuffer.mockImplementation((path) => { + if (path === imagePath1) { + return Promise.resolve({ dataUrl: base64DataUrl1, buffer: Buffer.from("test1") }) + } + if (path === imagePath2) { + return Promise.resolve({ dataUrl: base64DataUrl2, buffer: Buffer.from("test2") }) + } + return Promise.reject(new Error("File not found")) + }) + + // Mock server and tool exist + mockProviderRef.deref.mockReturnValue({ + getMcpHub: () => ({ + getAllServers: vi + .fn() + .mockReturnValue([ + { + name: "test_server", + tools: [{ name: "process_images", description: "Process multiple images" }], + }, + ]), + callTool: vi.fn().mockResolvedValue({ + content: [{ type: "text", text: "Images processed" }], + isError: false, + }), + }), + postMessageToWebview: vi.fn(), + }) + + const block: ToolUse = { + type: "tool_use", + name: "use_mcp_tool", + params: { + server_name: "test_server", + tool_name: "process_images", + arguments: JSON.stringify({ + images: [imagePath1, imagePath2], + metadata: { + primary_image: imagePath1, + thumbnail: imagePath2, + }, + text: "Some text that should not be converted", + }), + }, + partial: false, + } + + mockAskApproval.mockResolvedValue(true) + + // Execute + await useMcpToolTool.handle(mockTask as Task, block as any, { + askApproval: mockAskApproval, + handleError: mockHandleError, + pushToolResult: mockPushToolResult, + removeClosingTag: mockRemoveClosingTag, + toolProtocol: "xml", + }) + + // Verify the approval message contains converted base64 data URLs + const approvalCall = mockAskApproval.mock.calls[0] + const approvalMessage = JSON.parse(approvalCall[1]) + const args = JSON.parse(approvalMessage.arguments) + + expect(args.images[0]).toBe(base64DataUrl1) + expect(args.images[1]).toBe(base64DataUrl2) + expect(args.metadata.primary_image).toBe(base64DataUrl1) + expect(args.metadata.thumbnail).toBe(base64DataUrl2) + expect(args.text).toBe("Some text that should not be converted") + }) + + it("should skip conversion for non-image file paths", async () => { + // Setup + const textFilePath = "/path/to/document.txt" + + // Mock file exists but is not an image + mockFsStat.mockResolvedValue({ + isFile: () => true, + }) + + // Mock server and tool exist + mockProviderRef.deref.mockReturnValue({ + getMcpHub: () => ({ + getAllServers: vi + .fn() + .mockReturnValue([ + { name: "test_server", tools: [{ name: "process_file", description: "Process a file" }] }, + ]), + callTool: vi.fn().mockResolvedValue({ + content: [{ type: "text", text: "File processed" }], + isError: false, + }), + }), + postMessageToWebview: vi.fn(), + }) + + const block: ToolUse = { + type: "tool_use", + name: "use_mcp_tool", + params: { + server_name: "test_server", + tool_name: "process_file", + arguments: JSON.stringify({ file: textFilePath }), + }, + partial: false, + } + + mockAskApproval.mockResolvedValue(true) + + // Execute + await useMcpToolTool.handle(mockTask as Task, block as any, { + askApproval: mockAskApproval, + handleError: mockHandleError, + pushToolResult: mockPushToolResult, + removeClosingTag: mockRemoveClosingTag, + toolProtocol: "xml", + }) + + // Verify image conversion was NOT attempted + expect(mockReadImageAsDataUrlWithBuffer).not.toHaveBeenCalled() + + // Verify the original path is preserved + const approvalCall = mockAskApproval.mock.calls[0] + const approvalMessage = JSON.parse(approvalCall[1]) + const args = JSON.parse(approvalMessage.arguments) + expect(args.file).toBe(textFilePath) + }) + + it("should skip conversion for already base64 encoded images", async () => { + // Setup + const base64DataUrl = + "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg==" + + // Mock server and tool exist + mockProviderRef.deref.mockReturnValue({ + getMcpHub: () => ({ + getAllServers: vi + .fn() + .mockReturnValue([ + { + name: "test_server", + tools: [{ name: "process_image", description: "Process an image" }], + }, + ]), + callTool: vi.fn().mockResolvedValue({ + content: [{ type: "text", text: "Image processed" }], + isError: false, + }), + }), + postMessageToWebview: vi.fn(), + }) + + const block: ToolUse = { + type: "tool_use", + name: "use_mcp_tool", + params: { + server_name: "test_server", + tool_name: "process_image", + arguments: JSON.stringify({ image: base64DataUrl }), + }, + partial: false, + } + + mockAskApproval.mockResolvedValue(true) + + // Execute + await useMcpToolTool.handle(mockTask as Task, block as any, { + askApproval: mockAskApproval, + handleError: mockHandleError, + pushToolResult: mockPushToolResult, + removeClosingTag: mockRemoveClosingTag, + toolProtocol: "xml", + }) + + // Verify no file system operations were performed + expect(mockFsStat).not.toHaveBeenCalled() + expect(mockReadImageAsDataUrlWithBuffer).not.toHaveBeenCalled() + + // Verify the base64 data URL is preserved as-is + const approvalCall = mockAskApproval.mock.calls[0] + const approvalMessage = JSON.parse(approvalCall[1]) + const args = JSON.parse(approvalMessage.arguments) + expect(args.image).toBe(base64DataUrl) + }) + + it("should handle file read errors gracefully", async () => { + // Setup + const imagePath = "/path/to/nonexistent.png" + + // Mock file exists + mockFsStat.mockResolvedValue({ + isFile: () => true, + }) + + // Mock image reading failure + mockReadImageAsDataUrlWithBuffer.mockRejectedValue(new Error("File read error")) + + // Mock server and tool exist + mockProviderRef.deref.mockReturnValue({ + getMcpHub: () => ({ + getAllServers: vi + .fn() + .mockReturnValue([ + { + name: "test_server", + tools: [{ name: "process_image", description: "Process an image" }], + }, + ]), + callTool: vi.fn().mockResolvedValue({ + content: [{ type: "text", text: "Image processed" }], + isError: false, + }), + }), + postMessageToWebview: vi.fn(), + }) + + const block: ToolUse = { + type: "tool_use", + name: "use_mcp_tool", + params: { + server_name: "test_server", + tool_name: "process_image", + arguments: JSON.stringify({ image: imagePath }), + }, + partial: false, + } + + mockAskApproval.mockResolvedValue(true) + + // Execute + await useMcpToolTool.handle(mockTask as Task, block as any, { + askApproval: mockAskApproval, + handleError: mockHandleError, + pushToolResult: mockPushToolResult, + removeClosingTag: mockRemoveClosingTag, + toolProtocol: "xml", + }) + + // Verify the original path is preserved when conversion fails + const approvalCall = mockAskApproval.mock.calls[0] + const approvalMessage = JSON.parse(approvalCall[1]) + const args = JSON.parse(approvalMessage.arguments) + expect(args.image).toBe(imagePath) + }) + }) })