diff --git a/src/core/tools/UseMcpToolTool.ts b/src/core/tools/UseMcpToolTool.ts index 7546606fd7..6e819825af 100644 --- a/src/core/tools/UseMcpToolTool.ts +++ b/src/core/tools/UseMcpToolTool.ts @@ -319,9 +319,19 @@ export class UseMcpToolTool extends BaseTool<"use_mcp_tool"> { response: outputText || (images.length > 0 ? `[${images.length} image(s)]` : ""), }) - toolResultPretty = - (toolResult.isError ? "Error:\n" : "") + - (outputText || (images.length > 0 ? `[${images.length} image(s) received]` : "")) + // 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 (images.length > 0) { + const imageDataSection = images + .map((img, index) => `\n${img}\n`) + .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}` + resultText = resultText ? resultText + imageInfo : imageInfo.trim() + } + + toolResultPretty = (toolResult.isError ? "Error:\n" : "") + resultText } // Send completion status diff --git a/src/core/tools/__tests__/useMcpToolTool.spec.ts b/src/core/tools/__tests__/useMcpToolTool.spec.ts index 3a575e6218..f44266c763 100644 --- a/src/core/tools/__tests__/useMcpToolTool.spec.ts +++ b/src/core/tools/__tests__/useMcpToolTool.spec.ts @@ -618,14 +618,12 @@ describe("useMcpToolTool", () => { mockProviderRef.deref.mockReturnValue({ getMcpHub: () => ({ callTool: vi.fn().mockResolvedValue(mockToolResult), - getAllServers: vi - .fn() - .mockReturnValue([ - { - name: "figma-server", - tools: [{ name: "get_screenshot", description: "Get screenshot" }], - }, - ]), + getAllServers: vi.fn().mockReturnValue([ + { + name: "figma-server", + tools: [{ name: "get_screenshot", description: "Get screenshot" }], + }, + ]), }), postMessageToWebview: vi.fn(), }) @@ -637,9 +635,13 @@ describe("useMcpToolTool", () => { }) expect(mockTask.say).toHaveBeenCalledWith("mcp_server_request_started") - expect(mockTask.say).toHaveBeenCalledWith("mcp_server_response", "[1 image(s) received]", [ - "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJ", - ]) + expect(mockTask.say).toHaveBeenCalledWith( + "mcp_server_response", + expect.stringContaining( + "[1 image(s) received - data URLs provided below for use with save_image tool]", + ), + ["data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJ"], + ) expect(mockPushToolResult).toHaveBeenCalledWith(expect.stringContaining("with 1 image(s)")) }) @@ -693,9 +695,11 @@ describe("useMcpToolTool", () => { }) expect(mockTask.say).toHaveBeenCalledWith("mcp_server_request_started") - expect(mockTask.say).toHaveBeenCalledWith("mcp_server_response", "Node name: Button", [ - "data:image/png;base64,base64imagedata", - ]) + expect(mockTask.say).toHaveBeenCalledWith( + "mcp_server_response", + expect.stringContaining("Node name: Button"), + ["data:image/png;base64,base64imagedata"], + ) expect(mockPushToolResult).toHaveBeenCalledWith(expect.stringContaining("with 1 image(s)")) }) @@ -732,14 +736,12 @@ describe("useMcpToolTool", () => { mockProviderRef.deref.mockReturnValue({ getMcpHub: () => ({ callTool: vi.fn().mockResolvedValue(mockToolResult), - getAllServers: vi - .fn() - .mockReturnValue([ - { - name: "figma-server", - tools: [{ name: "get_screenshot", description: "Get screenshot" }], - }, - ]), + getAllServers: vi.fn().mockReturnValue([ + { + name: "figma-server", + tools: [{ name: "get_screenshot", description: "Get screenshot" }], + }, + ]), }), postMessageToWebview: vi.fn(), }) @@ -751,9 +753,13 @@ describe("useMcpToolTool", () => { }) // Should not double-prefix the data URL - expect(mockTask.say).toHaveBeenCalledWith("mcp_server_response", "[1 image(s) received]", [ - "data:image/jpeg;base64,/9j/4AAQSkZJRg==", - ]) + expect(mockTask.say).toHaveBeenCalledWith( + "mcp_server_response", + expect.stringContaining( + "[1 image(s) received - data URLs provided below for use with save_image tool]", + ), + ["data:image/jpeg;base64,/9j/4AAQSkZJRg=="], + ) }) it("should handle multiple images in response", async () => { @@ -794,14 +800,12 @@ describe("useMcpToolTool", () => { mockProviderRef.deref.mockReturnValue({ getMcpHub: () => ({ callTool: vi.fn().mockResolvedValue(mockToolResult), - getAllServers: vi - .fn() - .mockReturnValue([ - { - name: "figma-server", - tools: [{ name: "get_screenshots", description: "Get screenshots" }], - }, - ]), + getAllServers: vi.fn().mockReturnValue([ + { + name: "figma-server", + tools: [{ name: "get_screenshots", description: "Get screenshots" }], + }, + ]), }), postMessageToWebview: vi.fn(), }) @@ -812,10 +816,13 @@ describe("useMcpToolTool", () => { pushToolResult: mockPushToolResult, }) - expect(mockTask.say).toHaveBeenCalledWith("mcp_server_response", "[2 image(s) received]", [ - "data:image/png;base64,image1data", - "data:image/png;base64,image2data", - ]) + expect(mockTask.say).toHaveBeenCalledWith( + "mcp_server_response", + expect.stringContaining( + "[2 image(s) received - data URLs provided below for use with save_image tool]", + ), + ["data:image/png;base64,image1data", "data:image/png;base64,image2data"], + ) expect(mockPushToolResult).toHaveBeenCalledWith(expect.stringContaining("with 2 image(s)")) }) }) diff --git a/webview-ui/src/components/chat/ChatRow.tsx b/webview-ui/src/components/chat/ChatRow.tsx index c15d760cb5..66f9ba75f5 100644 --- a/webview-ui/src/components/chat/ChatRow.tsx +++ b/webview-ui/src/components/chat/ChatRow.tsx @@ -119,6 +119,7 @@ interface ChatRowProps { onFollowUpUnmount?: () => void isFollowUpAnswered?: boolean isFollowUpAutoApprovalPaused?: boolean + mcpResponseImages?: string[] editable?: boolean hasCheckpoint?: boolean } @@ -173,6 +174,7 @@ export const ChatRowContent = ({ onBatchFileResponse, isFollowUpAnswered, isFollowUpAutoApprovalPaused, + mcpResponseImages, }: ChatRowContentProps) => { const { t, i18n } = useTranslation() @@ -1627,7 +1629,7 @@ export const ChatRowContent = ({ server={server} useMcpServer={useMcpServer} alwaysAllowMcp={alwaysAllowMcp} - images={message.images} + images={mcpResponseImages ?? message.images} /> )} diff --git a/webview-ui/src/components/chat/ChatView.tsx b/webview-ui/src/components/chat/ChatView.tsx index 81f6cbebf6..fdf0ff648f 100644 --- a/webview-ui/src/components/chat/ChatView.tsx +++ b/webview-ui/src/components/chat/ChatView.tsx @@ -1327,6 +1327,15 @@ const ChatViewComponent: React.ForwardRefRenderFunction } + // For use_mcp_server ask messages, find the corresponding mcp_server_response to get images + let mcpResponseImages: string[] | undefined + if (messageOrGroup.type === "ask" && messageOrGroup.ask === "use_mcp_server") { + const mcpResponse = modifiedMessages.find( + (m) => m.ts > messageOrGroup.ts && m.say === "mcp_server_response", + ) + mcpResponseImages = mcpResponse?.images + } + // regular message return (