Fixes #5252: Handle images without data URI format in access_mcp_resource

roomote · roomote · commit 86e5ff3fa9e1 · 2025-06-30T11:14:01.000Z
- Modified accessMcpResourceTool.ts to detect and convert raw base64 data to proper data URI format
- Updated formatImagesIntoBlocks in responses.ts to handle both data URI and raw base64 formats
- Added comprehensive test suite covering standard data URI, raw base64, mixed formats, and different MIME types
- Maintains backward compatibility with existing data URI format
- Uses 'image/png' as default MIME type for raw base64 data
diff --git a/src/core/prompts/__tests__/responses-image-handling.spec.ts b/src/core/prompts/__tests__/responses-image-handling.spec.ts
@@ -0,0 +1,228 @@
+// npx vitest core/prompts/__tests__/responses-image-handling.spec.ts
+
+import { vi, describe, it, expect } from "vitest"
+import { formatResponse } from "../responses"
+
+// Mock VSCode dependencies
+vi.mock("vscode", () => {
+	const mockDisposable = { dispose: vi.fn() }
+	return {
+		workspace: {
+			createFileSystemWatcher: vi.fn(() => ({
+				onDidCreate: vi.fn(() => mockDisposable),
+				onDidChange: vi.fn(() => mockDisposable),
+				onDidDelete: vi.fn(() => mockDisposable),
+				dispose: vi.fn(),
+			})),
+		},
+		RelativePattern: vi.fn(),
+	}
+})
+
+// Mock fs dependencies
+vi.mock("../../../utils/fs", () => ({
+	fileExistsAtPath: vi.fn().mockResolvedValue(false),
+}))
+
+vi.mock("fs/promises", () => ({
+	readFile: vi.fn().mockResolvedValue(""),
+}))
+
+describe("Image Handling in formatResponse", () => {
+	describe("formatResponse.toolResult with images", () => {
+		it("should handle standard data URI format images", () => {
+			const text = "Here is an image:"
+			const images = [
+				"data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg==",
+			]
+
+			const result = formatResponse.toolResult(text, images)
+
+			// Should return an array with text and image blocks
+			expect(Array.isArray(result)).toBe(true)
+			const resultArray = result as any[]
+
+			// First block should be text
+			expect(resultArray[0]).toEqual({
+				type: "text",
+				text: "Here is an image:",
+			})
+
+			// Second block should be image with correct format
+			expect(resultArray[1]).toEqual({
+				type: "image",
+				source: {
+					type: "base64",
+					media_type: "image/png",
+					data: "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg==",
+				},
+			})
+		})
+
+		it("should handle raw base64 data without data URI prefix", () => {
+			const text = "Here is an image:"
+			const images = [
+				"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg==",
+			]
+
+			const result = formatResponse.toolResult(text, images)
+
+			// Should return an array with text and image blocks
+			expect(Array.isArray(result)).toBe(true)
+			const resultArray = result as any[]
+
+			// First block should be text
+			expect(resultArray[0]).toEqual({
+				type: "text",
+				text: "Here is an image:",
+			})
+
+			// Second block should be image with default PNG mime type
+			expect(resultArray[1]).toEqual({
+				type: "image",
+				source: {
+					type: "base64",
+					media_type: "image/png", // Default fallback
+					data: "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg==",
+				},
+			})
+		})
+
+		it("should handle different image MIME types in data URI", () => {
+			const text = "Here are different image types:"
+			const images = [
+				"data:image/jpeg;base64,/9j/4AAQSkZJRgABAQEAYABgAAD/2wBDAAEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQH/2wBDAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQH/wAARCAABAAEDASIAAhEBAxEB/8QAFQABAQAAAAAAAAAAAAAAAAAAAAv/xAAUEAEAAAAAAAAAAAAAAAAAAAAA/8QAFQEBAQAAAAAAAAAAAAAAAAAAAAX/xAAUEQEAAAAAAAAAAAAAAAAAAAAA/9oADAMBAAIRAxEAPwA/8A8A",
+				"data:image/gif;base64,R0lGODlhAQABAIAAAAAAAP///yH5BAEAAAAALAAAAAABAAEAAAIBRAA7",
+				"data:image/webp;base64,UklGRiIAAABXRUJQVlA4IBYAAAAwAQCdASoBAAEADsD+JaQAA3AAAAAA",
+			]
+
+			const result = formatResponse.toolResult(text, images)
+
+			// Should return an array with text and image blocks
+			expect(Array.isArray(result)).toBe(true)
+			const resultArray = result as any[]
+
+			// First block should be text
+			expect(resultArray[0]).toEqual({
+				type: "text",
+				text: "Here are different image types:",
+			})
+
+			// Should handle JPEG
+			expect(resultArray[1]).toEqual({
+				type: "image",
+				source: {
+					type: "base64",
+					media_type: "image/jpeg",
+					data: "/9j/4AAQSkZJRgABAQEAYABgAAD/2wBDAAEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQH/2wBDAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQH/wAARCAABAAEDASIAAhEBAxEB/8QAFQABAQAAAAAAAAAAAAAAAAAAAAv/xAAUEAEAAAAAAAAAAAAAAAAAAAAA/8QAFQEBAQAAAAAAAAAAAAAAAAAAAAX/xAAUEQEAAAAAAAAAAAAAAAAAAAAA/9oADAMBAAIRAxEAPwA/8A8A",
+				},
+			})
+
+			// Should handle GIF
+			expect(resultArray[2]).toEqual({
+				type: "image",
+				source: {
+					type: "base64",
+					media_type: "image/gif",
+					data: "R0lGODlhAQABAIAAAAAAAP///yH5BAEAAAAALAAAAAABAAEAAAIBRAA7",
+				},
+			})
+
+			// Should handle WebP
+			expect(resultArray[3]).toEqual({
+				type: "image",
+				source: {
+					type: "base64",
+					media_type: "image/webp",
+					data: "UklGRiIAAABXRUJQVlA4IBYAAAAwAQCdASoBAAEADsD+JaQAA3AAAAAA",
+				},
+			})
+		})
+
+		it("should handle mixed data URI and raw base64 images", () => {
+			const text = "Mixed image formats:"
+			const images = [
+				"data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg==",
+				"R0lGODlhAQABAIAAAAAAAP///yH5BAEAAAAALAAAAAABAAEAAAIBRAA7", // Raw base64 GIF
+			]
+
+			const result = formatResponse.toolResult(text, images)
+
+			// Should return an array with text and image blocks
+			expect(Array.isArray(result)).toBe(true)
+			const resultArray = result as any[]
+
+			// First image should preserve original PNG format
+			expect(resultArray[1]).toEqual({
+				type: "image",
+				source: {
+					type: "base64",
+					media_type: "image/png",
+					data: "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg==",
+				},
+			})
+
+			// Second image should default to PNG for raw base64
+			expect(resultArray[2]).toEqual({
+				type: "image",
+				source: {
+					type: "base64",
+					media_type: "image/png", // Default fallback
+					data: "R0lGODlhAQABAIAAAAAAAP///yH5BAEAAAAALAAAAAABAAEAAAIBRAA7",
+				},
+			})
+		})
+
+		it("should return just text when no images provided", () => {
+			const text = "Just text, no images"
+
+			const result = formatResponse.toolResult(text)
+
+			// Should return just the text string
+			expect(result).toBe("Just text, no images")
+		})
+
+		it("should return just text when empty images array provided", () => {
+			const text = "Just text, empty images array"
+			const images: string[] = []
+
+			const result = formatResponse.toolResult(text, images)
+
+			// Should return just the text string
+			expect(result).toBe("Just text, empty images array")
+		})
+	})
+
+	describe("formatResponse.imageBlocks", () => {
+		it("should handle undefined images", () => {
+			const result = formatResponse.imageBlocks(undefined)
+
+			expect(result).toEqual([])
+		})
+
+		it("should handle empty images array", () => {
+			const result = formatResponse.imageBlocks([])
+
+			expect(result).toEqual([])
+		})
+
+		it("should format raw base64 images correctly", () => {
+			const images = [
+				"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg==",
+			]
+
+			const result = formatResponse.imageBlocks(images)
+
+			expect(result).toEqual([
+				{
+					type: "image",
+					source: {
+						type: "base64",
+						media_type: "image/png",
+						data: "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg==",
+					},
+				},
+			])
+		})
+	})
+})
diff --git a/src/core/prompts/responses.ts b/src/core/prompts/responses.ts
@@ -179,13 +179,27 @@ Otherwise, if you have not completed the task and do not need additional informa
 const formatImagesIntoBlocks = (images?: string[]): Anthropic.ImageBlockParam[] => {
 	return images
 		? images.map((dataUrl) => {
-				// data:image/png;base64,base64string
-				const [rest, base64] = dataUrl.split(",")
-				const mimeType = rest.split(":")[1].split(";")[0]
-				return {
-					type: "image",
-					source: { type: "base64", media_type: mimeType, data: base64 },
-				} as Anthropic.ImageBlockParam
+				// Handle different image formats:
+				// 1. data:image/png;base64,base64string (standard data URI)
+				// 2. base64string (raw base64 data)
+				// 3. other formats that might be provided by MCP servers
+
+				if (dataUrl.startsWith("data:")) {
+					// Standard data URI format: data:image/png;base64,base64string
+					const [rest, base64] = dataUrl.split(",")
+					const mimeType = rest.split(":")[1].split(";")[0]
+					return {
+						type: "image",
+						source: { type: "base64", media_type: mimeType, data: base64 },
+					} as Anthropic.ImageBlockParam
+				} else {
+					// Assume it's raw base64 data, default to image/png
+					// This handles cases where MCP servers provide just the base64 string
+					return {
+						type: "image",
+						source: { type: "base64", media_type: "image/png", data: dataUrl },
+					} as Anthropic.ImageBlockParam
+				}
 			})
 		: []
 }
diff --git a/src/core/tools/accessMcpResourceTool.ts b/src/core/tools/accessMcpResourceTool.ts
@@ -73,7 +73,16 @@ export async function accessMcpResourceTool(
 
 			resourceResult?.contents.forEach((item) => {
 				if (item.mimeType?.startsWith("image") && item.blob) {
-					images.push(item.blob)
+					// Check if blob is already a data URI
+					if (item.blob.startsWith("data:")) {
+						// Already in data URI format, use as-is
+						images.push(item.blob)
+					} else {
+						// Assume it's raw base64 data, create proper data URI
+						const mimeType = item.mimeType || "image/png"
+						const dataUri = `data:${mimeType};base64,${item.blob}`
+						images.push(dataUri)
+					}
 				}
 			})
 
