feat: enhance image handling in MCP tool processing

**Changes:**
- Improved validation for base64 image data in `processToolContent` to handle invalid and non-string data gracefully.
- Added error handling to log warnings for corrupted images without interrupting processing.
- Updated tests to verify correct behavior when encountering invalid base64 data and non-string inputs.

**Files Modified:**
- `src/core/tools/useMcpToolTool.ts`
- `src/core/tools/__tests__/useMcpToolTool.spec.ts`
- `webview-ui/src/components/common/Thumbnails.tsx`

Author: shivangag

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

@@ -333,6 +333,135 @@ describe("useMcpToolTool", () => {
 			expect(mockPushToolResult).toHaveBeenCalledWith("Tool result:  (with 1 images)")
 		})
 
+		it("should handle corrupted base64 image data gracefully", async () => {
+			const block: ToolUse = {
+				type: "tool_use",
+				name: "use_mcp_tool",
+				params: {
+					server_name: "test_server",
+					tool_name: "test_tool",
+					arguments: '{"param": "value"}',
+				},
+				partial: false,
+			}
+
+			mockAskApproval.mockResolvedValue(true)
+
+			const mockToolResult = {
+				content: [
+					{ type: "text", text: "Generated content with images:" },
+					{
+						type: "image",
+						data: "invalid@base64@data", // Invalid base64 characters
+						mimeType: "image/png",
+					},
+					{
+						type: "image",
+						data: "", // Empty base64 data
+						mimeType: "image/png",
+					},
+					{
+						type: "image",
+						data: "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChAI9jU", // Valid base64
+						mimeType: "image/png",
+					},
+				],
+				isError: false,
+			}
+
+			mockProviderRef.deref.mockReturnValue({
+				getMcpHub: () => ({
+					callTool: vi.fn().mockResolvedValue(mockToolResult),
+				}),
+				postMessageToWebview: vi.fn(),
+			})
+
+			// Spy on console.warn to verify error logging
+			const consoleSpy = vi.spyOn(console, "warn").mockImplementation(() => {})
+
+			await useMcpToolTool(
+				mockTask as Task,
+				block,
+				mockAskApproval,
+				mockHandleError,
+				mockPushToolResult,
+				mockRemoveClosingTag,
+			)
+
+			// Should continue processing despite corrupted images
+			expect(mockTask.consecutiveMistakeCount).toBe(0)
+			expect(mockAskApproval).toHaveBeenCalled()
+			expect(mockTask.say).toHaveBeenCalledWith("mcp_server_request_started")
+
+			// Should only include the valid image, not the corrupted ones
+			expect(mockTask.say).toHaveBeenCalledWith("mcp_server_response", "Generated content with images:", [
+				"data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChAI9jU",
+			])
+			expect(mockPushToolResult).toHaveBeenCalledWith(
+				"Tool result: Generated content with images: (with 1 images)",
+			)
+
+			// Should log warnings for corrupted images
+			expect(consoleSpy).toHaveBeenCalledWith("Invalid MCP ImageContent: base64 data contains invalid characters")
+			expect(consoleSpy).toHaveBeenCalledWith("Invalid MCP ImageContent: base64 data is not a valid string")
+
+			consoleSpy.mockRestore()
+		})
+
+		it("should handle non-string base64 data", async () => {
+			const block: ToolUse = {
+				type: "tool_use",
+				name: "use_mcp_tool",
+				params: {
+					server_name: "test_server",
+					tool_name: "test_tool",
+					arguments: '{"param": "value"}',
+				},
+				partial: false,
+			}
+
+			mockAskApproval.mockResolvedValue(true)
+
+			const mockToolResult = {
+				content: [
+					{ type: "text", text: "Some text" },
+					{
+						type: "image",
+						data: 12345, // Non-string data
+						mimeType: "image/png",
+					},
+				],
+				isError: false,
+			}
+
+			mockProviderRef.deref.mockReturnValue({
+				getMcpHub: () => ({
+					callTool: vi.fn().mockResolvedValue(mockToolResult),
+				}),
+				postMessageToWebview: vi.fn(),
+			})
+
+			const consoleSpy = vi.spyOn(console, "warn").mockImplementation(() => {})
+
+			await useMcpToolTool(
+				mockTask as Task,
+				block,
+				mockAskApproval,
+				mockHandleError,
+				mockPushToolResult,
+				mockRemoveClosingTag,
+			)
+
+			// Should process text content normally
+			expect(mockTask.say).toHaveBeenCalledWith("mcp_server_response", "Some text", [])
+			expect(mockPushToolResult).toHaveBeenCalledWith("Tool result: Some text")
+
+			// Should log warning for invalid data type
+			expect(consoleSpy).toHaveBeenCalledWith("Invalid MCP ImageContent: base64 data is not a valid string")
+
+			consoleSpy.mockRestore()
+		})
+
 		it("should handle user rejection", async () => {
 			const block: ToolUse = {
 				type: "tool_use",

src/core/tools/useMcpToolTool.ts

@@ -207,11 +207,29 @@ function processToolContent(toolResult: any): { text: string; images: string[] }
 		if (item.type === "text") {
 			textParts.push(item.text)
 		} else if (item.type === "image") {
-			if (item.data && item.mimeType) {
+			if (item.mimeType && item.data !== undefined && item.data !== null) {
 				const validImageTypes = ["image/png", "image/jpeg", "image/gif", "image/webp"]
 				if (validImageTypes.includes(item.mimeType)) {
-					const dataUrl = `data:${item.mimeType};base64,${item.data}`
-					images.push(dataUrl)
+					try {
+						// Validate base64 data before constructing data URL
+						if (typeof item.data !== "string" || item.data.trim() === "") {
+							console.warn("Invalid MCP ImageContent: base64 data is not a valid string")
+							return
+						}
+
+						// Basic validation for base64 format
+						const base64Regex = /^[A-Za-z0-9+/]*={0,2}$/
+						if (!base64Regex.test(item.data.replace(/\s/g, ""))) {
+							console.warn("Invalid MCP ImageContent: base64 data contains invalid characters")
+							return
+						}
+
+						const dataUrl = `data:${item.mimeType};base64,${item.data}`
+						images.push(dataUrl)
+					} catch (error) {
+						console.warn("Failed to process MCP image content:", error)
+						// Continue processing other content instead of failing entirely
+					}
 				} else {
 					console.warn(`Unsupported image MIME type: ${item.mimeType}`)
 				}

webview-ui/src/components/common/Thumbnails.tsx

@@ -11,6 +11,7 @@ interface ThumbnailsProps {
 
 const Thumbnails = ({ images, style, setImages, onHeightChange }: ThumbnailsProps) => {
 	const [hoveredIndex, setHoveredIndex] = useState<number | null>(null)
+	const [failedImages, setFailedImages] = useState<Set<number>>(new Set())
 	const containerRef = useRef<HTMLDivElement>(null)
 	const { width } = useWindowSize()
 
@@ -24,12 +25,19 @@ const Thumbnails = ({ images, style, setImages, onHeightChange }: ThumbnailsProp
 			onHeightChange?.(height)
 		}
 		setHoveredIndex(null)
+		// Reset failed images when images change
+		setFailedImages(new Set())
 	}, [images, width, onHeightChange])
 
 	const handleDelete = (index: number) => {
 		setImages?.((prevImages) => prevImages.filter((_, i) => i !== index))
 	}
 
+	const handleImageError = (index: number) => {
+		setFailedImages((prev) => new Set(prev).add(index))
+		console.warn(`Failed to load image at index ${index}`)
+	}
+
 	const isDeletable = setImages !== undefined
 
 	const handleImageClick = (image: string) => {
@@ -53,18 +61,43 @@ const Thumbnails = ({ images, style, setImages, onHeightChange }: ThumbnailsProp
 					style={{ position: "relative" }}
 					onMouseEnter={() => setHoveredIndex(index)}
 					onMouseLeave={() => setHoveredIndex(null)}>
-					<img
-						src={image}
-						alt={`Thumbnail ${index + 1}`}
-						style={{
-							width: 34,
-							height: 34,
-							objectFit: "cover",
-							borderRadius: 4,
-							cursor: "pointer",
-						}}
-						onClick={() => handleImageClick(image)}
-					/>
+					{failedImages.has(index) ? (
+						<div
+							style={{
+								width: 34,
+								height: 34,
+								borderRadius: 4,
+								backgroundColor: "var(--vscode-input-background)",
+								border: "1px solid var(--vscode-input-border)",
+								display: "flex",
+								justifyContent: "center",
+								alignItems: "center",
+								cursor: "pointer",
+							}}
+							onClick={() => handleImageClick(image)}
+							title="Failed to load image">
+							<span
+								className="codicon codicon-file-media"
+								style={{
+									color: "var(--vscode-descriptionForeground)",
+									fontSize: 16,
+								}}></span>
+						</div>
+					) : (
+						<img
+							src={image}
+							alt={`Thumbnail ${index + 1}`}
+							style={{
+								width: 34,
+								height: 34,
+								objectFit: "cover",
+								borderRadius: 4,
+								cursor: "pointer",
+							}}
+							onClick={() => handleImageClick(image)}
+							onError={() => handleImageError(index)}
+						/>
+					)}
 					{isDeletable && hoveredIndex === index && (
 						<div
 							onClick={() => handleDelete(index)}