fix: address self-review comments for image mentions feature

- Maintain backward compatibility for parseMentions() return type
- Add error handling for corrupted images in processImageFile()
- Use i18n for memory limit error message
- Improve test clarity with @/path syntax
- Add documentation for array flattening logic
- Add test coverage for SVG, WebP, and AVIF formats

src/core/mentions/__tests__/imageMentions.spec.ts

@@ -140,11 +140,15 @@ describe("Image Mentions", () => {
 				20,
 			)
 
-			expect(result.images).toHaveLength(2)
-			expect(result.images[0]).toBe(mockImageDataUrl1)
-			expect(result.images[1]).toBe(mockImageDataUrl2)
-			expect(result.text).toContain("'image1.png' (see below for image)")
-			expect(result.text).toContain("'image2.jpg' (see below for image)")
+			// Type guard for TypeScript
+			expect(typeof result).toBe("object")
+			if (typeof result === "object") {
+				expect(result.images).toHaveLength(2)
+				expect(result.images[0]).toBe(mockImageDataUrl1)
+				expect(result.images[1]).toBe(mockImageDataUrl2)
+				expect(result.text).toContain("'image1.png' (see below for image)")
+				expect(result.text).toContain("'image2.jpg' (see below for image)")
+			}
 		})
 
 		it("should handle image size limit exceeded", async () => {
@@ -177,8 +181,11 @@ describe("Image Mentions", () => {
 				20,
 			)
 
-			expect(result.images).toHaveLength(0)
-			expect(result.text).toContain("Image file is too large (10 MB). Maximum allowed size is 5 MB.")
+			// When validation fails, we still get an object but with no images
+			expect(typeof result).toBe("string")
+			if (typeof result === "string") {
+				expect(result).toContain("Image file is too large (10 MB). Maximum allowed size is 5 MB.")
+			}
 		})
 
 		it("should handle model that doesn't support images", async () => {
@@ -210,8 +217,11 @@ describe("Image Mentions", () => {
 				20,
 			)
 
-			expect(result.images).toHaveLength(0)
-			expect(result.text).toContain("Image file detected but current model does not support images")
+			// When model doesn't support images, result should be a string
+			expect(typeof result).toBe("string")
+			if (typeof result === "string") {
+				expect(result).toContain("Image file detected but current model does not support images")
+			}
 		})
 
 		it("should handle mixed content with images and regular files", async () => {
@@ -263,12 +273,16 @@ describe("Image Mentions", () => {
 				20,
 			)
 
-			expect(result.images).toHaveLength(1)
-			expect(result.images[0]).toBe(mockImageDataUrl)
-			expect(result.text).toContain("'image.png' (see below for image)")
-			// The script.js file will have an error because we're not fully mocking the file system
-			// but that's okay for this test - we're mainly testing that images and non-images are handled differently
-			expect(result.text).toContain("script.js")
+			// Type guard for TypeScript
+			expect(typeof result).toBe("object")
+			if (typeof result === "object") {
+				expect(result.images).toHaveLength(1)
+				expect(result.images[0]).toBe(mockImageDataUrl)
+				expect(result.text).toContain("'image.png' (see below for image)")
+				// The script.js file will have an error because we're not fully mocking the file system
+				// but that's okay for this test - we're mainly testing that images and non-images are handled differently
+				expect(result.text).toContain("script.js")
+			}
 		})
 
 		it("should respect .rooignore for image files", async () => {
@@ -299,8 +313,11 @@ describe("Image Mentions", () => {
 				20,
 			)
 
-			expect(result.images).toHaveLength(0)
-			expect(result.text).toContain("(Image ignored-image.png is ignored by .rooignore)")
+			// When image is ignored, result should be a string
+			expect(typeof result).toBe("string")
+			if (typeof result === "string") {
+				expect(result).toContain("(Image ignored-image.png is ignored by .rooignore)")
+			}
 		})
 
 		it("should handle total memory limit for multiple images", async () => {
@@ -348,8 +365,159 @@ describe("Image Mentions", () => {
 				20, // maxTotalImageSize
 			)
 
-			expect(result.images).toHaveLength(1)
-			expect(result.text).toContain("Image skipped to avoid size limit")
+			// Type guard for TypeScript
+			expect(typeof result).toBe("object")
+			if (typeof result === "object") {
+				expect(result.images).toHaveLength(1)
+				expect(result.text).toContain("Image skipped to avoid size limit")
+			}
+		})
+
+		it("should handle SVG image format", async () => {
+			const mockSvgDataUrl =
+				"data:image/svg+xml;base64,PHN2ZyB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciPjwvc3ZnPg=="
+
+			vi.mocked(imageHelpers.isSupportedImageFormat).mockImplementation((ext) => ext === ".svg")
+			vi.mocked(imageHelpers.validateImageForProcessing).mockResolvedValue({
+				isValid: true,
+				sizeInMB: 0.1,
+			})
+
+			vi.mocked(imageHelpers.processImageFile).mockResolvedValue({
+				dataUrl: mockSvgDataUrl,
+				buffer: Buffer.from('<svg xmlns="http://www.w3.org/2000/svg"></svg>'),
+				sizeInKB: 100,
+				sizeInMB: 0.1,
+				notice: "Image (100 KB)",
+			})
+
+			vi.mocked(fs.stat).mockResolvedValue({
+				isFile: () => true,
+				isDirectory: () => false,
+				size: 102400,
+			} as any)
+
+			const result = await parseMentions(
+				"Check @/logo.svg",
+				"/workspace",
+				mockUrlContentFetcher,
+				undefined,
+				undefined,
+				true,
+				true,
+				50,
+				undefined,
+				true,
+				5,
+				20,
+			)
+
+			// Type guard for TypeScript
+			expect(typeof result).toBe("object")
+			if (typeof result === "object") {
+				expect(result).toHaveProperty("images")
+				expect(result.images).toHaveLength(1)
+				expect(result.images[0]).toBe(mockSvgDataUrl)
+				expect(result.text).toContain("'logo.svg' (see below for image)")
+			}
+		})
+
+		it("should handle WebP image format", async () => {
+			const mockWebpDataUrl =
+				"data:image/webp;base64,UklGRiQAAABXRUJQVlA4IBgAAAAwAQCdASoBAAEAAQAcJaQAA3AA/v3AgAA="
+
+			vi.mocked(imageHelpers.isSupportedImageFormat).mockImplementation((ext) => ext === ".webp")
+			vi.mocked(imageHelpers.validateImageForProcessing).mockResolvedValue({
+				isValid: true,
+				sizeInMB: 0.2,
+			})
+
+			vi.mocked(imageHelpers.processImageFile).mockResolvedValue({
+				dataUrl: mockWebpDataUrl,
+				buffer: Buffer.from("webp data"),
+				sizeInKB: 200,
+				sizeInMB: 0.2,
+				notice: "Image (200 KB)",
+			})
+
+			vi.mocked(fs.stat).mockResolvedValue({
+				isFile: () => true,
+				isDirectory: () => false,
+				size: 204800,
+			} as any)
+
+			const result = await parseMentions(
+				"Check @/photo.webp",
+				"/workspace",
+				mockUrlContentFetcher,
+				undefined,
+				undefined,
+				true,
+				true,
+				50,
+				undefined,
+				true,
+				5,
+				20,
+			)
+
+			// Type guard for TypeScript
+			expect(typeof result).toBe("object")
+			if (typeof result === "object") {
+				expect(result).toHaveProperty("images")
+				expect(result.images).toHaveLength(1)
+				expect(result.images[0]).toBe(mockWebpDataUrl)
+				expect(result.text).toContain("'photo.webp' (see below for image)")
+			}
+		})
+
+		it("should handle AVIF image format", async () => {
+			const mockAvifDataUrl =
+				"data:image/avif;base64,AAAAIGZ0eXBhdmlmAAAAAGF2aWZtaWYxbWlhZk1BMUIAAADybWV0YQAAAAAAAAAoaGRscgAAAAAAAAAAcGljdAAAAAAAAAAAAAAAAGxpYmF2aWYAAAAADnBpdG0AAAAAAAEAAAAeaWxvYwAAAABEAAABAAEAAAABAAABGgAAAB0AAAAoaWluZgAAAAAAAQAAABppbmZlAgAAAAABAABhdjAxQ29sb3IAAAAAamlwcnAAAABLaXBjbwAAABRpc3BlAAAAAAAAAAEAAAABAAAAEHBpeGkAAAAAAwgICAAAAAxhdjFDgQ0MAAAAABNjb2xybmNseAACAAIABoAAAAAXaXBtYQAAAAAAAAABAAEEAQKDBAAAACVtZGF0EgAKCBgANogQEAwgMg=="
+
+			vi.mocked(imageHelpers.isSupportedImageFormat).mockImplementation((ext) => ext === ".avif")
+			vi.mocked(imageHelpers.validateImageForProcessing).mockResolvedValue({
+				isValid: true,
+				sizeInMB: 0.3,
+			})
+
+			vi.mocked(imageHelpers.processImageFile).mockResolvedValue({
+				dataUrl: mockAvifDataUrl,
+				buffer: Buffer.from("avif data"),
+				sizeInKB: 300,
+				sizeInMB: 0.3,
+				notice: "Image (300 KB)",
+			})
+
+			vi.mocked(fs.stat).mockResolvedValue({
+				isFile: () => true,
+				isDirectory: () => false,
+				size: 307200,
+			} as any)
+
+			const result = await parseMentions(
+				"Check @/modern.avif",
+				"/workspace",
+				mockUrlContentFetcher,
+				undefined,
+				undefined,
+				true,
+				true,
+				50,
+				undefined,
+				true,
+				5,
+				20,
+			)
+
+			// Type guard for TypeScript
+			expect(typeof result).toBe("object")
+			if (typeof result === "object") {
+				expect(result).toHaveProperty("images")
+				expect(result.images).toHaveLength(1)
+				expect(result.images[0]).toBe(mockAvifDataUrl)
+				expect(result.text).toContain("'modern.avif' (see below for image)")
+			}
 		})
 	})
 })

src/core/mentions/index.ts

@@ -99,7 +99,7 @@ export async function parseMentions(
 	supportsImages: boolean = false,
 	maxImageFileSize: number = DEFAULT_MAX_IMAGE_FILE_SIZE_MB,
 	maxTotalImageSize: number = DEFAULT_MAX_TOTAL_IMAGE_SIZE_MB,
-): Promise<{ text: string; images: string[] }> {
+): Promise<string | { text: string; images: string[] }> {
 	const mentions: Set<string> = new Set()
 	const validCommands: Map<string, Command> = new Map()
 	const imageDataUrls: string[] = []
@@ -327,6 +327,11 @@ export async function parseMentions(
 		}
 	}
 
+	// Maintain backward compatibility: if no images were found, return just the string
+	// Otherwise, return the new format with text and images
+	if (imageDataUrls.length === 0) {
+		return parsedText
+	}
 	return { text: parsedText, images: imageDataUrls }
 }
 

src/core/mentions/processUserContentMentions.ts

@@ -69,6 +69,15 @@ export async function processUserContentMentions({
 						maxTotalImageSize,
 					)
 
+					// Handle backward compatibility - result can be string or object
+					if (typeof result === "string") {
+						// No images found, just return the text block with updated text
+						return {
+							...block,
+							text: result,
+						}
+					}
+
 					// If there are images, we need to add them as separate image blocks
 					const blocks: Anthropic.Messages.ContentBlockParam[] = [
 						{
@@ -116,10 +125,13 @@ export async function processUserContentMentions({
 							maxTotalImageSize,
 						)
 
+						// Handle backward compatibility - result can be string or object
+						const textContent = typeof result === "string" ? result : result.text
+
 						// For tool_result, we can only return text content, not images
 						return {
 							...block,
-							content: result.text,
+							content: textContent,
 						}
 					}
 
@@ -143,10 +155,13 @@ export async function processUserContentMentions({
 									maxTotalImageSize,
 								)
 
+								// Handle backward compatibility - result can be string or object
+								const textContent = typeof result === "string" ? result : result.text
+
 								// For tool_result content blocks, we can only return text
 								return {
 									...contentBlock,
-									text: result.text,
+									text: textContent,
 								}
 							}
 
@@ -163,7 +178,10 @@ export async function processUserContentMentions({
 			return block
 		}),
 	).then((results) => {
-		// Flatten any arrays that were returned (when images were added)
+		// Flatten any arrays that were returned (when images were added).
+		// This is necessary because when we process image mentions, we return an array
+		// containing both the text block and separate image blocks. The flat() method
+		// ensures all blocks are at the same level in the final array.
 		return results.flat()
 	})
 }

src/core/tools/helpers/imageHelpers.ts

@@ -73,14 +73,25 @@ export interface ImageProcessingResult {
  * Reads an image file and returns both the data URL and buffer
  */
 export async function readImageAsDataUrlWithBuffer(filePath: string): Promise<{ dataUrl: string; buffer: Buffer }> {
-	const fileBuffer = await fs.readFile(filePath)
-	const base64 = fileBuffer.toString("base64")
-	const ext = path.extname(filePath).toLowerCase()
+	try {
+		const fileBuffer = await fs.readFile(filePath)
 
-	const mimeType = IMAGE_MIME_TYPES[ext] || "image/png"
-	const dataUrl = `data:${mimeType};base64,${base64}`
+		// Basic validation to check if the buffer is not empty
+		if (!fileBuffer || fileBuffer.length === 0) {
+			throw new Error("Image file is empty or corrupted")
+		}
 
-	return { dataUrl, buffer: fileBuffer }
+		const base64 = fileBuffer.toString("base64")
+		const ext = path.extname(filePath).toLowerCase()
+
+		const mimeType = IMAGE_MIME_TYPES[ext] || "image/png"
+		const dataUrl = `data:${mimeType};base64,${base64}`
+
+		return { dataUrl, buffer: fileBuffer }
+	} catch (error) {
+		// Re-throw with more context
+		throw new Error(`Failed to read image file: ${error.message}`)
+	}
 }
 
 /**
@@ -133,7 +144,11 @@ export async function validateImageForProcessing(
 		return {
 			isValid: false,
 			reason: "memory_limit",
-			notice: `Image skipped to avoid size limit (${maxTotalImageSize}MB). Current: ${currentMemoryFormatted} + this file: ${fileMemoryFormatted}. Try fewer or smaller images.`,
+			notice: t("tools:readFile.imageTotalSizeExceeded", {
+				maxTotal: maxTotalImageSize,
+				current: currentMemoryFormatted,
+				fileSize: fileMemoryFormatted,
+			}),
 			sizeInMB: imageSizeInMB,
 		}
 	}
@@ -148,18 +163,47 @@ export async function validateImageForProcessing(
  * Processes an image file and returns the result
  */
 export async function processImageFile(fullPath: string): Promise<ImageProcessingResult> {
-	const imageStats = await fs.stat(fullPath)
-	const { dataUrl, buffer } = await readImageAsDataUrlWithBuffer(fullPath)
-	const imageSizeInKB = Math.round(imageStats.size / 1024)
-	const imageSizeInMB = imageStats.size / (1024 * 1024)
-	const noticeText = t("tools:readFile.imageWithSize", { size: imageSizeInKB })
+	try {
+		const imageStats = await fs.stat(fullPath)
 
-	return {
-		dataUrl,
-		buffer,
-		sizeInKB: imageSizeInKB,
-		sizeInMB: imageSizeInMB,
-		notice: noticeText,
+		// Validate file exists and is not empty
+		if (!imageStats.isFile()) {
+			throw new Error("Path does not point to a file")
+		}
+
+		if (imageStats.size === 0) {
+			throw new Error("Image file is empty")
+		}
+
+		const { dataUrl, buffer } = await readImageAsDataUrlWithBuffer(fullPath)
+
+		// Additional validation on the buffer
+		if (!buffer || buffer.length === 0) {
+			throw new Error("Failed to read image data")
+		}
+
+		const imageSizeInKB = Math.round(imageStats.size / 1024)
+		const imageSizeInMB = imageStats.size / (1024 * 1024)
+		const noticeText = t("tools:readFile.imageWithSize", { size: imageSizeInKB })
+
+		return {
+			dataUrl,
+			buffer,
+			sizeInKB: imageSizeInKB,
+			sizeInMB: imageSizeInMB,
+			notice: noticeText,
+		}
+	} catch (error) {
+		// Provide more context about the error
+		if (error.code === "ENOENT") {
+			throw new Error(`Image file not found: ${fullPath}`)
+		} else if (error.code === "EACCES") {
+			throw new Error(`Permission denied accessing image file: ${fullPath}`)
+		} else if (error.message.includes("corrupted")) {
+			throw new Error(`Image file appears to be corrupted: ${fullPath}`)
+		}
+		// Re-throw with original message if not a known error
+		throw error
 	}
 }
 

src/i18n/locales/en/tools.json

@@ -4,7 +4,8 @@
 		"definitionsOnly": " (definitions only)",
 		"maxLines": " (max {{max}} lines)",
 		"imageTooLarge": "Image file is too large ({{size}} MB). The maximum allowed size is {{max}} MB.",
-		"imageWithSize": "Image file ({{size}} KB)"
+		"imageWithSize": "Image file ({{size}} KB)",
+		"imageTotalSizeExceeded": "Image skipped to avoid size limit ({{maxTotal}}MB). Current: {{current}} + this file: {{fileSize}}. Try fewer or smaller images."
 	},
 	"toolRepetitionLimitReached": "Roo appears to be stuck in a loop, attempting the same action ({{toolName}}) repeatedly. This might indicate a problem with its current strategy. Consider rephrasing the task, providing more specific instructions, or guiding it towards a different approach.",
 	"codebaseSearch": {