feat: add image support for MCP tool responses

- Handle image content from MCP tool responses
- Save images to local file system with proper naming
- Display images in UI alongside text responses
- Add comprehensive tests for image handling

Fixes #9817

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 path from "path"
+import * as fs from "fs/promises"
+import { getWorkspacePath } from "../../utils/path"
 
 interface UseMcpToolParams {
 	server_name: string
@@ -268,24 +271,76 @@ export class UseMcpToolTool extends BaseTool<"use_mcp_tool"> {
 		})
 	}
 
-	private processToolContent(toolResult: any): string {
+	private processToolContent(toolResult: any): { text: string; images: string[] } {
 		if (!toolResult?.content || toolResult.content.length === 0) {
-			return ""
+			return { text: "", images: [] }
 		}
 
-		return toolResult.content
-			.map((item: any) => {
-				if (item.type === "text") {
-					return item.text
+		const textParts: string[] = []
+		const images: string[] = []
+
+		toolResult.content.forEach((item: any) => {
+			if (item.type === "text") {
+				textParts.push(item.text)
+			} else if (item.type === "resource") {
+				const { blob: _, ...rest } = item.resource
+				textParts.push(JSON.stringify(rest, null, 2))
+			} else if (item.type === "image" && item.data && item.mimeType) {
+				// Handle image content from MCP tool response
+				let imageDataUrl: string
+				if (item.data.startsWith("data:")) {
+					imageDataUrl = item.data
+				} else {
+					// Construct data URL from base64 data and mimeType
+					imageDataUrl = `data:${item.mimeType};base64,${item.data}`
 				}
-				if (item.type === "resource") {
-					const { blob: _, ...rest } = item.resource
-					return JSON.stringify(rest, null, 2)
-				}
-				return ""
-			})
-			.filter(Boolean)
-			.join("\n\n")
+				images.push(imageDataUrl)
+			}
+		})
+
+		return {
+			text: textParts.filter(Boolean).join("\n\n"),
+			images,
+		}
+	}
+
+	private async saveImageToFile(dataUrl: string, index: number): Promise<string | null> {
+		try {
+			// Parse the data URL
+			const matches = dataUrl.match(/^data:image\/([a-zA-Z]+);base64,(.+)$/)
+			if (!matches) {
+				console.error("Invalid image data URL format")
+				return null
+			}
+
+			const [, format, base64Data] = matches
+			const imageBuffer = Buffer.from(base64Data, "base64")
+
+			// Get workspace path or use temp directory
+			const workspacePath = getWorkspacePath()
+			const saveDir = workspacePath || require("os").tmpdir()
+
+			// Create images directory if it doesn't exist
+			const imagesDir = path.join(saveDir, "mcp-images")
+			await fs.mkdir(imagesDir, { recursive: true })
+
+			// Generate filename with timestamp
+			const timestamp = new Date().toISOString().replace(/[:.]/g, "-").slice(0, -5)
+			const filename = `mcp-image-${timestamp}-${index}.${format}`
+			const filePath = path.join(imagesDir, filename)
+
+			// Write the image file
+			await fs.writeFile(filePath, imageBuffer)
+
+			// Return relative path if in workspace, absolute path otherwise
+			if (workspacePath) {
+				return path.relative(workspacePath, filePath)
+			}
+			return filePath
+		} catch (error) {
+			console.error("Error saving image to file:", error)
+			return null
+		}
 	}
 
 	private async executeToolAndProcessResult(
@@ -309,18 +364,41 @@ export class UseMcpToolTool extends BaseTool<"use_mcp_tool"> {
 		const toolResult = await task.providerRef.deref()?.getMcpHub()?.callTool(serverName, toolName, parsedArguments)
 
 		let toolResultPretty = "(No response)"
+		let images: string[] = []
+		let savedImagePaths: string[] = []
 
 		if (toolResult) {
-			const outputText = this.processToolContent(toolResult)
+			const processedContent = this.processToolContent(toolResult)
+			const outputText = processedContent.text
+			images = processedContent.images
+
+			// Save images to local file system
+			if (images.length > 0) {
+				for (let i = 0; i < images.length; i++) {
+					const savedPath = await this.saveImageToFile(images[i], i + 1)
+					if (savedPath) {
+						savedImagePaths.push(savedPath)
+					}
+				}
+			}
+
+			if (outputText || images.length > 0) {
+				// Include saved image paths in the output text
+				let responseText = outputText
+				if (savedImagePaths.length > 0) {
+					const imagePathsText = savedImagePaths
+						.map((path, index) => `Image ${index + 1} saved to: ${path}`)
+						.join("\n")
+					responseText = outputText ? `${outputText}\n\n${imagePathsText}` : imagePathsText
+				}
 
-			if (outputText) {
 				await this.sendExecutionStatus(task, {
 					executionId,
 					status: "output",
-					response: outputText,
+					response: responseText,
 				})
 
-				toolResultPretty = (toolResult.isError ? "Error:\n" : "") + outputText
+				toolResultPretty = (toolResult.isError ? "Error:\n" : "") + responseText
 			}
 
 			// Send completion status
@@ -339,8 +417,9 @@ export class UseMcpToolTool extends BaseTool<"use_mcp_tool"> {
 			})
 		}
 
-		await task.say("mcp_server_response", toolResultPretty)
-		pushToolResult(formatResponse.toolResult(toolResultPretty))
+		// Pass images to task.say and formatResponse.toolResult for display in UI
+		await task.say("mcp_server_response", toolResultPretty, images)
+		pushToolResult(formatResponse.toolResult(toolResultPretty, images))
 	}
 }
 

src/core/tools/__tests__/useMcpToolTool.spec.ts

@@ -3,11 +3,17 @@
 import { useMcpToolTool } from "../UseMcpToolTool"
 import { Task } from "../../task/Task"
 import { ToolUse } from "../../../shared/tools"
+import * as fs from "fs/promises"
 
 // Mock dependencies
 vi.mock("../../prompts/responses", () => ({
 	formatResponse: {
-		toolResult: vi.fn((result: string) => `Tool result: ${result}`),
+		toolResult: vi.fn((result: string, images?: string[]) => {
+			if (images && images.length > 0) {
+				return `Tool result: ${result} [with ${images.length} image(s)]`
+			}
+			return `Tool result: ${result}`
+		}),
 		toolError: vi.fn((error: string) => `Tool error: ${error}`),
 		invalidMcpToolArgumentError: vi.fn((server: string, tool: string) => `Invalid args for ${server}:${tool}`),
 		unknownMcpToolError: vi.fn((server: string, tool: string, availableTools: string[]) => {
@@ -21,6 +27,17 @@ vi.mock("../../prompts/responses", () => ({
 	},
 }))
 
+// Mock fs/promises for image saving tests
+vi.mock("fs/promises", () => ({
+	mkdir: vi.fn(),
+	writeFile: vi.fn(),
+}))
+
+// Mock getWorkspacePath
+vi.mock("../../../utils/path", () => ({
+	getWorkspacePath: vi.fn(() => "/workspace"),
+}))
+
 vi.mock("../../../i18n", () => ({
 	t: vi.fn((key: string, params?: any) => {
 		if (key === "mcp:errors.invalidJsonArgument" && params?.toolName) {
@@ -234,7 +251,7 @@ describe("useMcpToolTool", () => {
 			expect(mockTask.consecutiveMistakeCount).toBe(0)
 			expect(mockAskApproval).toHaveBeenCalled()
 			expect(mockTask.say).toHaveBeenCalledWith("mcp_server_request_started")
-			expect(mockTask.say).toHaveBeenCalledWith("mcp_server_response", "Tool executed successfully")
+			expect(mockTask.say).toHaveBeenCalledWith("mcp_server_response", "Tool executed successfully", [])
 			expect(mockPushToolResult).toHaveBeenCalledWith("Tool result: Tool executed successfully")
 		})
 
@@ -448,7 +465,7 @@ describe("useMcpToolTool", () => {
 			expect(mockTask.consecutiveMistakeCount).toBe(0)
 			expect(mockTask.recordToolError).not.toHaveBeenCalled()
 			expect(mockTask.say).toHaveBeenCalledWith("mcp_server_request_started")
-			expect(mockTask.say).toHaveBeenCalledWith("mcp_server_response", "Tool executed successfully")
+			expect(mockTask.say).toHaveBeenCalledWith("mcp_server_response", "Tool executed successfully", [])
 		})
 
 		it("should reject unknown server names with available servers listed", async () => {
@@ -537,4 +554,300 @@ describe("useMcpToolTool", () => {
 			expect(mockAskApproval).not.toHaveBeenCalled()
 		})
 	})
+
+	describe("image handling", () => {
+		beforeEach(() => {
+			// Setup fs mocks
+			vi.mocked(fs).mkdir.mockResolvedValue(undefined as any)
+			vi.mocked(fs).writeFile.mockResolvedValue(undefined as any)
+		})
+
+		it("should handle tool response with image content", async () => {
+			const block: ToolUse = {
+				type: "tool_use",
+				name: "use_mcp_tool",
+				params: {
+					server_name: "test_server",
+					tool_name: "image_tool",
+					arguments: '{"prompt": "generate image"}',
+				},
+				partial: false,
+			}
+
+			mockAskApproval.mockResolvedValue(true)
+
+			const mockToolResult = {
+				content: [
+					{ type: "text", text: "Image generated successfully" },
+					{
+						type: "image",
+						data: "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg==",
+						mimeType: "image/png",
+					},
+				],
+				isError: false,
+			}
+
+			mockProviderRef.deref.mockReturnValue({
+				getMcpHub: () => ({
+					callTool: vi.fn().mockResolvedValue(mockToolResult),
+				}),
+				postMessageToWebview: vi.fn(),
+			})
+
+			await useMcpToolTool.handle(mockTask as Task, block as any, {
+				askApproval: mockAskApproval,
+				handleError: mockHandleError,
+				pushToolResult: mockPushToolResult,
+				removeClosingTag: mockRemoveClosingTag,
+				toolProtocol: "xml",
+			})
+
+			// Verify image was saved to file system
+			expect(vi.mocked(fs).mkdir).toHaveBeenCalledWith(expect.stringContaining("mcp-images"), { recursive: true })
+			expect(vi.mocked(fs).writeFile).toHaveBeenCalledWith(expect.stringContaining(".png"), expect.any(Buffer))
+
+			// Verify task.say was called with images
+			expect(mockTask.say).toHaveBeenCalledWith(
+				"mcp_server_response",
+				expect.stringContaining("Image generated successfully"),
+				expect.arrayContaining([expect.stringContaining("data:image/png;base64,")]),
+			)
+
+			// Verify pushToolResult was called with images
+			expect(mockPushToolResult).toHaveBeenCalledWith(expect.stringContaining("[with 1 image(s)]"))
+		})
+
+		it("should handle multiple images in tool response", async () => {
+			const block: ToolUse = {
+				type: "tool_use",
+				name: "use_mcp_tool",
+				params: {
+					server_name: "test_server",
+					tool_name: "multi_image_tool",
+					arguments: "{}",
+				},
+				partial: false,
+			}
+
+			mockAskApproval.mockResolvedValue(true)
+
+			const mockToolResult = {
+				content: [
+					{
+						type: "image",
+						data: "image1base64data",
+						mimeType: "image/jpeg",
+					},
+					{
+						type: "image",
+						data: "image2base64data",
+						mimeType: "image/png",
+					},
+					{ type: "text", text: "Generated 2 images" },
+				],
+				isError: false,
+			}
+
+			mockProviderRef.deref.mockReturnValue({
+				getMcpHub: () => ({
+					callTool: vi.fn().mockResolvedValue(mockToolResult),
+				}),
+				postMessageToWebview: vi.fn(),
+			})
+
+			// Clear previous mock calls from other tests
+			vi.mocked(fs).writeFile.mockClear()
+			vi.mocked(fs).mkdir.mockClear()
+
+			await useMcpToolTool.handle(mockTask as Task, block as any, {
+				askApproval: mockAskApproval,
+				handleError: mockHandleError,
+				pushToolResult: mockPushToolResult,
+				removeClosingTag: mockRemoveClosingTag,
+				toolProtocol: "xml",
+			})
+
+			// Verify both images were saved
+			expect(vi.mocked(fs).writeFile).toHaveBeenCalledTimes(2)
+			expect(vi.mocked(fs).writeFile).toHaveBeenNthCalledWith(
+				1,
+				expect.stringContaining(".jpeg"),
+				expect.any(Buffer),
+			)
+			expect(vi.mocked(fs).writeFile).toHaveBeenNthCalledWith(
+				2,
+				expect.stringContaining(".png"),
+				expect.any(Buffer),
+			)
+
+			// Verify task.say was called with both images
+			expect(mockTask.say).toHaveBeenCalledWith(
+				"mcp_server_response",
+				expect.any(String),
+				expect.arrayContaining([
+					expect.stringContaining("data:image/jpeg;base64,"),
+					expect.stringContaining("data:image/png;base64,"),
+				]),
+			)
+
+			// Verify pushToolResult indicates multiple images
+			expect(mockPushToolResult).toHaveBeenCalledWith(expect.stringContaining("[with 2 image(s)]"))
+		})
+
+		it("should handle image with data URL format", async () => {
+			const block: ToolUse = {
+				type: "tool_use",
+				name: "use_mcp_tool",
+				params: {
+					server_name: "test_server",
+					tool_name: "image_tool",
+					arguments: "{}",
+				},
+				partial: false,
+			}
+
+			mockAskApproval.mockResolvedValue(true)
+
+			const mockToolResult = {
+				content: [
+					{
+						type: "image",
+						data: "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg==",
+						mimeType: "image/png",
+					},
+				],
+				isError: false,
+			}
+
+			mockProviderRef.deref.mockReturnValue({
+				getMcpHub: () => ({
+					callTool: vi.fn().mockResolvedValue(mockToolResult),
+				}),
+				postMessageToWebview: vi.fn(),
+			})
+
+			await useMcpToolTool.handle(mockTask as Task, block as any, {
+				askApproval: mockAskApproval,
+				handleError: mockHandleError,
+				pushToolResult: mockPushToolResult,
+				removeClosingTag: mockRemoveClosingTag,
+				toolProtocol: "xml",
+			})
+
+			// Verify image was processed correctly even with data URL format
+			expect(vi.mocked(fs).writeFile).toHaveBeenCalled()
+			expect(mockTask.say).toHaveBeenCalledWith(
+				"mcp_server_response",
+				expect.any(String),
+				expect.arrayContaining([expect.stringContaining("data:image/png;base64,")]),
+			)
+		})
+
+		it("should handle mixed content (text, resource, and image)", async () => {
+			const block: ToolUse = {
+				type: "tool_use",
+				name: "use_mcp_tool",
+				params: {
+					server_name: "test_server",
+					tool_name: "mixed_content_tool",
+					arguments: "{}",
+				},
+				partial: false,
+			}
+
+			mockAskApproval.mockResolvedValue(true)
+
+			const mockToolResult = {
+				content: [
+					{ type: "text", text: "Processing complete" },
+					{
+						type: "resource",
+						resource: {
+							uri: "file://test.txt",
+							mimeType: "text/plain",
+						},
+					},
+					{
+						type: "image",
+						data: "testImageData",
+						mimeType: "image/jpeg",
+					},
+				],
+				isError: false,
+			}
+
+			mockProviderRef.deref.mockReturnValue({
+				getMcpHub: () => ({
+					callTool: vi.fn().mockResolvedValue(mockToolResult),
+				}),
+				postMessageToWebview: vi.fn(),
+			})
+
+			await useMcpToolTool.handle(mockTask as Task, block as any, {
+				askApproval: mockAskApproval,
+				handleError: mockHandleError,
+				pushToolResult: mockPushToolResult,
+				removeClosingTag: mockRemoveClosingTag,
+				toolProtocol: "xml",
+			})
+
+			// Verify all content types were processed
+			expect(mockTask.say).toHaveBeenCalledWith(
+				"mcp_server_response",
+				expect.stringContaining("Processing complete"),
+				expect.arrayContaining([expect.stringContaining("data:image/jpeg;base64,")]),
+			)
+			expect(vi.mocked(fs).writeFile).toHaveBeenCalledWith(expect.stringContaining(".jpeg"), expect.any(Buffer))
+		})
+
+		it("should handle tool response with only images (no text)", async () => {
+			const block: ToolUse = {
+				type: "tool_use",
+				name: "use_mcp_tool",
+				params: {
+					server_name: "test_server",
+					tool_name: "image_only_tool",
+					arguments: "{}",
+				},
+				partial: false,
+			}
+
+			mockAskApproval.mockResolvedValue(true)
+
+			const mockToolResult = {
+				content: [
+					{
+						type: "image",
+						data: "imageData",
+						mimeType: "image/png",
+					},
+				],
+				isError: false,
+			}
+
+			mockProviderRef.deref.mockReturnValue({
+				getMcpHub: () => ({
+					callTool: vi.fn().mockResolvedValue(mockToolResult),
+				}),
+				postMessageToWebview: vi.fn(),
+			})
+
+			await useMcpToolTool.handle(mockTask as Task, block as any, {
+				askApproval: mockAskApproval,
+				handleError: mockHandleError,
+				pushToolResult: mockPushToolResult,
+				removeClosingTag: mockRemoveClosingTag,
+				toolProtocol: "xml",
+			})
+
+			// Verify image was saved and paths were included in response
+			expect(vi.mocked(fs).writeFile).toHaveBeenCalled()
+			expect(mockTask.say).toHaveBeenCalledWith(
+				"mcp_server_response",
+				expect.stringContaining("Image 1 saved to:"),
+				expect.arrayContaining([expect.stringContaining("data:image/png;base64,")]),
+			)
+		})
+	})
 })