fix: display images returned by MCP tools in chat and forward embedded resource images to the model

Author: taltas

src/core/tools/UseMcpToolTool.ts

@@ -290,7 +290,14 @@ export class UseMcpToolTool extends BaseTool<"use_mcp_tool"> {
 					return item.text
 				}
 				if (item.type === "resource") {
-					const { blob: _, ...rest } = item.resource
+					const { blob, ...rest } = item.resource
+					if (blob && item.resource.mimeType?.startsWith("image")) {
+						if (blob.startsWith("data:")) {
+							images.push(blob)
+						} else {
+							images.push(`data:${item.resource.mimeType};base64,${blob}`)
+						}
+					}
 					return JSON.stringify(rest, null, 2)
 				}
 				if (item.type === "image") {

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

@@ -859,6 +859,126 @@ describe("useMcpToolTool", () => {
 			])
 		})
 
+		it("should extract image data from embedded resource blobs", async () => {
+			const block: ToolUse<"use_mcp_tool"> = {
+				type: "tool_use",
+				name: "use_mcp_tool",
+				params: {
+					server_name: "godot-server",
+					tool_name: "game_screenshot",
+					arguments: "{}",
+				},
+				nativeArgs: {
+					server_name: "godot-server",
+					tool_name: "game_screenshot",
+					arguments: {},
+				},
+				partial: false,
+			}
+
+			mockAskApproval.mockResolvedValue(true)
+
+			const mockToolResult = {
+				content: [
+					{ type: "text", text: "Screenshot captured: 1152x648" },
+					{
+						type: "resource",
+						resource: {
+							uri: "godot://screenshot/latest",
+							mimeType: "image/png",
+							blob: "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJ",
+						},
+					},
+				],
+				isError: false,
+			}
+
+			mockProviderRef.deref.mockReturnValue({
+				getMcpHub: () => ({
+					callTool: vi.fn().mockResolvedValue(mockToolResult),
+					getAllServers: vi.fn().mockReturnValue([
+						{
+							name: "godot-server",
+							tools: [{ name: "game_screenshot", description: "Capture screenshot" }],
+						},
+					]),
+				}),
+				postMessageToWebview: vi.fn(),
+			})
+
+			await useMcpToolTool.handle(mockTask as Task, block, {
+				askApproval: mockAskApproval,
+				handleError: mockHandleError,
+				pushToolResult: mockPushToolResult,
+			})
+
+			expect(mockTask.say).toHaveBeenCalledWith(
+				"mcp_server_response",
+				expect.stringContaining("Screenshot captured: 1152x648"),
+				["data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJ"],
+			)
+			expect(mockPushToolResult).toHaveBeenCalledWith(expect.stringContaining("with 1 image(s)"))
+		})
+
+		it("should not extract non-image resource blobs", async () => {
+			const block: ToolUse<"use_mcp_tool"> = {
+				type: "tool_use",
+				name: "use_mcp_tool",
+				params: {
+					server_name: "godot-server",
+					tool_name: "read_scene",
+					arguments: "{}",
+				},
+				nativeArgs: {
+					server_name: "godot-server",
+					tool_name: "read_scene",
+					arguments: {},
+				},
+				partial: false,
+			}
+
+			mockAskApproval.mockResolvedValue(true)
+
+			const mockToolResult = {
+				content: [
+					{
+						type: "resource",
+						resource: {
+							uri: "godot://scene/main",
+							mimeType: "text/plain",
+							blob: "c2NlbmUgZGF0YQ==",
+						},
+					},
+				],
+				isError: false,
+			}
+
+			mockProviderRef.deref.mockReturnValue({
+				getMcpHub: () => ({
+					callTool: vi.fn().mockResolvedValue(mockToolResult),
+					getAllServers: vi.fn().mockReturnValue([
+						{
+							name: "godot-server",
+							tools: [{ name: "read_scene", description: "Read scene" }],
+						},
+					]),
+				}),
+				postMessageToWebview: vi.fn(),
+			})
+
+			await useMcpToolTool.handle(mockTask as Task, block, {
+				askApproval: mockAskApproval,
+				handleError: mockHandleError,
+				pushToolResult: mockPushToolResult,
+			})
+
+			expect(mockTask.say).toHaveBeenCalledWith(
+				"mcp_server_response",
+				expect.stringContaining("godot://scene/main"),
+				[],
+			)
+		})
+
 		it("should handle multiple images in response", async () => {
 			const block: ToolUse = {
 				type: "tool_use",

webview-ui/src/components/chat/ChatRow.tsx

@@ -1577,6 +1577,19 @@ export const ChatRowContent = ({
 						/>
 					)
 				}
+				case "mcp_server_response":
+					return (
+						<div style={{ paddingTop: 10 }}>
+							<Markdown markdown={message.text} partial={message.partial} />
+							{message.images && message.images.length > 0 && (
+								<div style={{ marginTop: "10px" }}>
+									{message.images.map((image, index) => (
+										<ImageBlock key={index} imageData={image} />
+									))}
+								</div>
+							)}
+						</div>
+					)
 				default:
 					return (
 						<>
@@ -1588,6 +1601,13 @@ export const ChatRowContent = ({
 							)}
 							<div style={{ paddingTop: 10 }}>
 								<Markdown markdown={message.text} partial={message.partial} />
+								{message.images && message.images.length > 0 && (
+									<div style={{ marginTop: "10px" }}>
+										{message.images.map((image, index) => (
+											<ImageBlock key={index} imageData={image} />
+										))}
+									</div>
+								)}
 							</div>
 						</>
 					)

webview-ui/src/components/chat/__tests__/ChatRow.mcp-server-response.spec.tsx

@@ -0,0 +1,90 @@
+import React from "react"
+
+import { render, screen } from "@/utils/test-utils"
+import { QueryClient, QueryClientProvider } from "@tanstack/react-query"
+import { ExtensionStateContextProvider } from "@src/context/ExtensionStateContext"
+import { ChatRowContent } from "../ChatRow"
+
+// Mock i18n
+vi.mock("react-i18next", () => ({
+	useTranslation: () => ({
+		t: (key: string) => key,
+		i18n: {
+			exists: () => false,
+		},
+	}),
+	Trans: ({ children }: { children?: React.ReactNode }) => <>{children}</>,
+	initReactI18next: { type: "3rdParty", init: () => {} },
+}))
+
+const queryClient = new QueryClient()
+
+function renderChatRow(message: any) {
+	return render(
+		<ExtensionStateContextProvider>
+			<QueryClientProvider client={queryClient}>
+				<ChatRowContent
+					message={message}
+					isExpanded={false}
+					isLast={false}
+					isStreaming={false}
+					onToggleExpand={() => {}}
+					onSuggestionClick={() => {}}
+					onBatchFileResponse={() => {}}
+					onFollowUpUnmount={() => {}}
+					isFollowUpAnswered={false}
+				/>
+			</QueryClientProvider>
+		</ExtensionStateContextProvider>,
+	)
+}
+
+describe("ChatRow - mcp_server_response", () => {
+	it("renders images attached to the MCP server response", () => {
+		const message: any = {
+			type: "say",
+			say: "mcp_server_response",
+			ts: Date.now(),
+			partial: false,
+			text: "Screenshot captured: 1152x648",
+			images: ["data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJ"],
+		}
+
+		renderChatRow(message)
+
+		const img = screen.getByRole("img")
+		expect(img).toHaveAttribute("src", "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJ")
+	})
+
+	it("renders multiple images attached to the MCP server response", () => {
+		const message: any = {
+			type: "say",
+			say: "mcp_server_response",
+			ts: Date.now(),
+			partial: false,
+			text: "[2 image(s) received]",
+			images: ["data:image/png;base64,image1data", "data:image/png;base64,image2data"],
+		}
+
+		renderChatRow(message)
+
+		const imgs = screen.getAllByRole("img")
+		expect(imgs).toHaveLength(2)
+		expect(imgs[0]).toHaveAttribute("src", "data:image/png;base64,image1data")
+		expect(imgs[1]).toHaveAttribute("src", "data:image/png;base64,image2data")
+	})
+
+	it("renders only the text when the MCP server response has no images", () => {
+		const message: any = {
+			type: "say",
+			say: "mcp_server_response",
+			ts: Date.now(),
+			partial: false,
+			text: "Plain text result",
+		}
+
+		renderChatRow(message)
+
+		expect(screen.queryByRole("img")).toBeNull()
+	})
+})