diff --git a/src/core/tools/UseMcpToolTool.ts b/src/core/tools/UseMcpToolTool.ts index b44b96054e..20c2a6b402 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 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 { + 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)) } } diff --git a/src/core/tools/__tests__/useMcpToolTool.spec.ts b/src/core/tools/__tests__/useMcpToolTool.spec.ts index 130047ae15..9d1d4ffab1 100644 --- a/src/core/tools/__tests__/useMcpToolTool.spec.ts +++ b/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,")]), + ) + }) + }) })