diff --git a/src/core/tools/UseMcpToolTool.ts b/src/core/tools/UseMcpToolTool.ts index 66a66b0dae..d605cc2440 100644 --- a/src/core/tools/UseMcpToolTool.ts +++ b/src/core/tools/UseMcpToolTool.ts @@ -329,11 +329,11 @@ export class UseMcpToolTool extends BaseTool<"use_mcp_tool"> { // File paths enable efficient persistence for future context, while raw base64 data allows // the model to analyze images in the current turn. if (images.length > 0) { - const savedImagePaths = await this.saveImagesToTempStorage(task, images, serverName, toolName) - const imagePathsSection = savedImagePaths + const savedImages = await this.saveImagesToTempStorage(task, images, serverName, toolName) + const imagePathsSection = savedImages .map( - (imgPath, index) => - `\n ${imgPath}\n ${images[index]}\n`, + (savedImage, index) => + `\n ${savedImage.path}\n ${images[savedImage.originalIndex]}\n`, ) .join("\n\n") const imageInfo = `\n\n[${images.length} image(s) received and saved to temporary storage. Use save_image tool with source_path (preferred) or data to save to your desired location.]\n\n${imagePathsSection}` @@ -364,17 +364,20 @@ export class UseMcpToolTool extends BaseTool<"use_mcp_tool"> { } /** - * Save images to task-specific temp storage and return file paths. + * Save images to task-specific temp storage and return file paths with original indices. * This allows passing file paths to the LLM instead of raw base64 data, * which prevents data corruption and reduces token costs. + * The originalIndex tracks each saved image's position in the source array so that + * callers can correctly pair paths with the original data URLs even when some + * images are skipped (e.g. unsupported formats). */ private async saveImagesToTempStorage( task: Task, images: string[], serverName: string, toolName: string, - ): Promise { - const savedPaths: string[] = [] + ): Promise<{ path: string; originalIndex: number }[]> { + const savedImages: { path: string; originalIndex: number }[] = [] try { const provider = task.providerRef.deref() @@ -406,15 +409,15 @@ export class UseMcpToolTool extends BaseTool<"use_mcp_tool"> { const imageBuffer = Buffer.from(data, "base64") await fs.writeFile(filePath, imageBuffer) - savedPaths.push(filePath) + savedImages.push({ path: filePath, originalIndex: i }) } } } 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 empty array on error - the LLM will see the error and handle accordingly } - return savedPaths + return savedImages } /** @@ -425,8 +428,8 @@ export class UseMcpToolTool extends BaseTool<"use_mcp_tool"> { images: string[], serverName: string, toolName: string, - ): Promise { - const savedPaths: string[] = [] + ): Promise<{ path: string; originalIndex: number }[]> { + const savedImages: { path: string; originalIndex: number }[] = [] try { const tempDir = path.join(task.cwd, ".roo", "temp", "mcp_images") @@ -445,14 +448,14 @@ export class UseMcpToolTool extends BaseTool<"use_mcp_tool"> { const imageBuffer = Buffer.from(data, "base64") await fs.writeFile(filePath, imageBuffer) - savedPaths.push(filePath) + savedImages.push({ path: filePath, originalIndex: i }) } } } catch (error) { console.error("Error saving images to fallback location:", error) } - return savedPaths + return savedImages } /** diff --git a/src/core/tools/__tests__/useMcpToolTool.spec.ts b/src/core/tools/__tests__/useMcpToolTool.spec.ts index 402b4454d6..fbf883b1c1 100644 --- a/src/core/tools/__tests__/useMcpToolTool.spec.ts +++ b/src/core/tools/__tests__/useMcpToolTool.spec.ts @@ -867,5 +867,88 @@ describe("useMcpToolTool", () => { ) expect(mockPushToolResult).toHaveBeenCalledWith(expect.stringContaining("with 2 image(s)")) }) + + it("should correctly pair saved paths with original images when unsupported formats are skipped", async () => { + const block: ToolUse = { + type: "tool_use", + name: "use_mcp_tool", + params: { + server_name: "figma-server", + tool_name: "get_screenshots", + arguments: '{"nodeIds": ["1", "2", "3"]}', + }, + nativeArgs: { + server_name: "figma-server", + tool_name: "get_screenshots", + arguments: { nodeIds: ["1", "2", "3"] }, + }, + partial: false, + } + + mockAskApproval.mockResolvedValue(true) + + // Mix of supported (png) and unsupported (tiff) image formats. + // processToolContent pushes all three into images[], but + // saveImagesToTempStorage skips the tiff because parseImageDataUrl + // doesn't match it. + const mockToolResult = { + content: [ + { + type: "image", + mimeType: "image/png", + data: "firstPngData", + }, + { + type: "image", + mimeType: "image/tiff", + data: "tiffDataSkipped", + }, + { + type: "image", + mimeType: "image/png", + data: "secondPngData", + }, + ], + isError: false, + } + + mockProviderRef.deref.mockReturnValue({ + getMcpHub: () => ({ + callTool: vi.fn().mockResolvedValue(mockToolResult), + getAllServers: vi.fn().mockReturnValue([ + { + name: "figma-server", + tools: [{ name: "get_screenshots", description: "Get screenshots" }], + }, + ]), + }), + postMessageToWebview: vi.fn(), + }) + + await useMcpToolTool.handle(mockTask as Task, block as any, { + askApproval: mockAskApproval, + handleError: mockHandleError, + pushToolResult: mockPushToolResult, + }) + + // The tiff image ends up in images[] as "data:image/tiff;base64,tiffDataSkipped" + // but parseImageDataUrl won't match it, so only 2 paths are saved. + // Thanks to originalIndex tracking, each pairs with the + // correct value even though the middle image was skipped. + const sayCall = vi.mocked(mockTask.say!).mock.calls.find( + (call: any[]) => + call[0] === "mcp_server_response" && typeof call[1] === "string" && call[1].includes(""), + ) + + expect(sayCall).toBeDefined() + const responseText = sayCall![1] as string + + // image_1 should pair with the first png (index 0), not the tiff + expect(responseText).toContain("data:image/png;base64,firstPngData") + // image_2 should pair with the second png (index 2), not the tiff + expect(responseText).toContain("data:image/png;base64,secondPngData") + // The tiff data URL should NOT appear in any tag + expect(responseText).not.toContain("tiffDataSkipped") + }) }) })