mirror of
https://github.com/RooVetGit/Roo-Code.git
synced 2026-09-06 08:18:39 +00:00
Fixes #5169: Add image reading support to read_file tool
- Extended FileResult interface to include imageData and mimeType fields - Added helper functions for image detection, MIME type mapping, and base64 conversion - Implemented size validation (10MB limit) to prevent memory issues - Added support for PNG, JPG, JPEG, WebP, GIF, BMP, ICO, TIFF, SVG formats - Created comprehensive test suite covering all image formats and edge cases - Images are now converted to base64 and included in XML output with metadata
This commit is contained in:
Roo Code
parent
3a8ba27615
commit
8bdd350555
2 changed files with 374 additions and 0 deletions
281
src/core/tools/__tests__/readFileTool.image.spec.ts
Normal file
281
src/core/tools/__tests__/readFileTool.image.spec.ts
Normal file
|
|
@ -0,0 +1,281 @@
|
|||
import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"
|
||||
import path from "path"
|
||||
import fs from "fs/promises"
|
||||
import { isBinaryFile } from "isbinaryfile"
|
||||
import { readFileTool } from "../readFileTool"
|
||||
import { Task } from "../../task/Task"
|
||||
import { ToolUse } from "../../../shared/tools"
|
||||
import { countFileLines } from "../../../integrations/misc/line-counter"
|
||||
import { extractTextFromFile, getSupportedBinaryFormats, addLineNumbers } from "../../../integrations/misc/extract-text"
|
||||
import { parseSourceCodeDefinitionsForFile } from "../../../services/tree-sitter"
|
||||
|
||||
// Mock dependencies
|
||||
vi.mock("fs/promises", () => ({
|
||||
default: {
|
||||
stat: vi.fn(),
|
||||
readFile: vi.fn(),
|
||||
},
|
||||
}))
|
||||
vi.mock("isbinaryfile")
|
||||
vi.mock("../../task/Task")
|
||||
vi.mock("../../../integrations/misc/line-counter")
|
||||
vi.mock("../../../integrations/misc/extract-text")
|
||||
vi.mock("../../../services/tree-sitter")
|
||||
|
||||
// Create mock functions
|
||||
const getSupportedBinaryFormatsMock = vi.fn(() => [".pdf", ".docx", ".ipynb", ".xlsx"])
|
||||
const extractTextFromFileMock = vi.fn()
|
||||
const addLineNumbersMock = vi.fn()
|
||||
|
||||
// Mock RooIgnoreController to handle vscode import
|
||||
vi.mock("../../ignore/RooIgnoreController", () => ({
|
||||
RooIgnoreController: class {
|
||||
initialize() {
|
||||
return Promise.resolve()
|
||||
}
|
||||
validateAccess() {
|
||||
return true
|
||||
}
|
||||
},
|
||||
}))
|
||||
|
||||
// Mock other dependencies
|
||||
vi.mock("../../../utils/fs", () => ({
|
||||
fileExistsAtPath: vi.fn().mockReturnValue(true),
|
||||
}))
|
||||
|
||||
const mockFs = vi.mocked(fs)
|
||||
const mockIsBinaryFile = vi.mocked(isBinaryFile)
|
||||
const mockedCountFileLines = vi.mocked(countFileLines)
|
||||
const mockedExtractTextFromFile = vi.mocked(extractTextFromFile)
|
||||
const mockedGetSupportedBinaryFormats = vi.mocked(getSupportedBinaryFormats)
|
||||
const mockedAddLineNumbers = vi.mocked(addLineNumbers)
|
||||
const mockedParseSourceCodeDefinitionsForFile = vi.mocked(parseSourceCodeDefinitionsForFile)
|
||||
|
||||
describe("readFileTool - Image Support", () => {
|
||||
let mockTask: any
|
||||
let mockAskApproval: any
|
||||
let mockHandleError: any
|
||||
let mockPushToolResult: any
|
||||
let mockRemoveClosingTag: any
|
||||
let toolResults: string[]
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks()
|
||||
toolResults = []
|
||||
|
||||
// Setup mocks
|
||||
mockedGetSupportedBinaryFormats.mockReturnValue([".pdf", ".docx", ".ipynb", ".xlsx"])
|
||||
mockedCountFileLines.mockResolvedValue(0)
|
||||
mockedExtractTextFromFile.mockResolvedValue("")
|
||||
mockedAddLineNumbers.mockImplementation((text) => text)
|
||||
mockedParseSourceCodeDefinitionsForFile.mockResolvedValue("")
|
||||
|
||||
mockTask = {
|
||||
cwd: "/test/workspace",
|
||||
rooIgnoreController: {
|
||||
validateAccess: vi.fn().mockReturnValue(true),
|
||||
},
|
||||
fileContextTracker: {
|
||||
trackFileContext: vi.fn(),
|
||||
},
|
||||
providerRef: {
|
||||
deref: vi.fn().mockReturnValue({
|
||||
getState: vi.fn().mockResolvedValue({ maxReadFileLine: -1 }),
|
||||
}),
|
||||
},
|
||||
ask: vi.fn().mockResolvedValue({ response: "yesButtonClicked" }),
|
||||
}
|
||||
|
||||
mockAskApproval = vi.fn()
|
||||
mockHandleError = vi.fn()
|
||||
mockPushToolResult = vi.fn((result: string) => {
|
||||
toolResults.push(result)
|
||||
})
|
||||
mockRemoveClosingTag = vi.fn()
|
||||
})
|
||||
|
||||
afterEach(() => {
|
||||
vi.resetAllMocks()
|
||||
})
|
||||
|
||||
it("should read PNG image file as base64", async () => {
|
||||
const imagePath = "test-image.png"
|
||||
const imageBuffer = Buffer.from("fake-png-data")
|
||||
const expectedBase64 = imageBuffer.toString("base64")
|
||||
|
||||
// Mock file operations
|
||||
mockIsBinaryFile.mockResolvedValue(true)
|
||||
mockFs.stat.mockResolvedValue({ size: 1024 } as any)
|
||||
mockFs.readFile.mockResolvedValue(imageBuffer)
|
||||
|
||||
const block: ToolUse = {
|
||||
type: "tool_use",
|
||||
name: "read_file",
|
||||
params: {
|
||||
args: `<file><path>${imagePath}</path></file>`,
|
||||
},
|
||||
partial: false,
|
||||
}
|
||||
|
||||
await readFileTool(mockTask, block, mockAskApproval, mockHandleError, mockPushToolResult, mockRemoveClosingTag)
|
||||
|
||||
expect(toolResults).toHaveLength(1)
|
||||
expect(toolResults[0]).toContain(`<image_data mime_type="image/png" size="1024">${expectedBase64}</image_data>`)
|
||||
expect(mockTask.fileContextTracker.trackFileContext).toHaveBeenCalledWith(imagePath, "read_tool")
|
||||
})
|
||||
|
||||
it("should read JPEG image file as base64", async () => {
|
||||
const imagePath = "test-image.jpg"
|
||||
const imageBuffer = Buffer.from("fake-jpeg-data")
|
||||
const expectedBase64 = imageBuffer.toString("base64")
|
||||
|
||||
// Mock file operations
|
||||
mockIsBinaryFile.mockResolvedValue(true)
|
||||
mockFs.stat.mockResolvedValue({ size: 2048 } as any)
|
||||
mockFs.readFile.mockResolvedValue(imageBuffer)
|
||||
|
||||
const block: ToolUse = {
|
||||
type: "tool_use",
|
||||
name: "read_file",
|
||||
params: {
|
||||
args: `<file><path>${imagePath}</path></file>`,
|
||||
},
|
||||
partial: false,
|
||||
}
|
||||
|
||||
await readFileTool(mockTask, block, mockAskApproval, mockHandleError, mockPushToolResult, mockRemoveClosingTag)
|
||||
|
||||
expect(toolResults).toHaveLength(1)
|
||||
expect(toolResults[0]).toContain(
|
||||
`<image_data mime_type="image/jpeg" size="2048">${expectedBase64}</image_data>`,
|
||||
)
|
||||
})
|
||||
|
||||
it("should handle multiple image files", async () => {
|
||||
const imagePaths = ["image1.png", "image2.jpg"]
|
||||
const imageBuffers = [Buffer.from("fake-png-data"), Buffer.from("fake-jpeg-data")]
|
||||
|
||||
// Mock file operations
|
||||
mockIsBinaryFile.mockResolvedValue(true)
|
||||
mockFs.stat.mockResolvedValueOnce({ size: 1024 } as any).mockResolvedValueOnce({ size: 2048 } as any)
|
||||
mockFs.readFile.mockResolvedValueOnce(imageBuffers[0]).mockResolvedValueOnce(imageBuffers[1])
|
||||
|
||||
const block: ToolUse = {
|
||||
type: "tool_use",
|
||||
name: "read_file",
|
||||
params: {
|
||||
args: `<file><path>${imagePaths[0]}</path></file><file><path>${imagePaths[1]}</path></file>`,
|
||||
},
|
||||
partial: false,
|
||||
}
|
||||
|
||||
await readFileTool(mockTask, block, mockAskApproval, mockHandleError, mockPushToolResult, mockRemoveClosingTag)
|
||||
|
||||
expect(toolResults).toHaveLength(1)
|
||||
expect(toolResults[0]).toContain(`<image_data mime_type="image/png"`)
|
||||
expect(toolResults[0]).toContain(`<image_data mime_type="image/jpeg"`)
|
||||
expect(toolResults[0]).toContain(imageBuffers[0].toString("base64"))
|
||||
expect(toolResults[0]).toContain(imageBuffers[1].toString("base64"))
|
||||
})
|
||||
|
||||
it("should reject image files that are too large", async () => {
|
||||
const imagePath = "large-image.png"
|
||||
const largeSize = 15 * 1024 * 1024 // 15MB (exceeds 10MB limit)
|
||||
|
||||
// Mock file operations
|
||||
mockIsBinaryFile.mockResolvedValue(true)
|
||||
mockFs.stat.mockResolvedValue({ size: largeSize } as any)
|
||||
|
||||
const block: ToolUse = {
|
||||
type: "tool_use",
|
||||
name: "read_file",
|
||||
params: {
|
||||
args: `<file><path>${imagePath}</path></file>`,
|
||||
},
|
||||
partial: false,
|
||||
}
|
||||
|
||||
await readFileTool(mockTask, block, mockAskApproval, mockHandleError, mockPushToolResult, mockRemoveClosingTag)
|
||||
|
||||
expect(toolResults).toHaveLength(1)
|
||||
expect(toolResults[0]).toContain("<error>Error reading image file:")
|
||||
expect(toolResults[0]).toContain("Image file is too large")
|
||||
expect(mockHandleError).toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it("should handle unsupported image formats as binary files", async () => {
|
||||
const imagePath = "test-image.tga" // Unsupported format
|
||||
|
||||
// Mock file operations
|
||||
mockIsBinaryFile.mockResolvedValue(true)
|
||||
|
||||
const block: ToolUse = {
|
||||
type: "tool_use",
|
||||
name: "read_file",
|
||||
params: {
|
||||
args: `<file><path>${imagePath}</path></file>`,
|
||||
},
|
||||
partial: false,
|
||||
}
|
||||
|
||||
await readFileTool(mockTask, block, mockAskApproval, mockHandleError, mockPushToolResult, mockRemoveClosingTag)
|
||||
|
||||
expect(toolResults).toHaveLength(1)
|
||||
expect(toolResults[0]).toContain("<notice>Binary file</notice>")
|
||||
})
|
||||
|
||||
it("should support WebP images", async () => {
|
||||
const imagePath = "test-image.webp"
|
||||
const imageBuffer = Buffer.from("fake-webp-data")
|
||||
const expectedBase64 = imageBuffer.toString("base64")
|
||||
|
||||
// Mock file operations
|
||||
mockIsBinaryFile.mockResolvedValue(true)
|
||||
mockFs.stat.mockResolvedValue({ size: 1024 } as any)
|
||||
mockFs.readFile.mockResolvedValue(imageBuffer)
|
||||
|
||||
const block: ToolUse = {
|
||||
type: "tool_use",
|
||||
name: "read_file",
|
||||
params: {
|
||||
args: `<file><path>${imagePath}</path></file>`,
|
||||
},
|
||||
partial: false,
|
||||
}
|
||||
|
||||
await readFileTool(mockTask, block, mockAskApproval, mockHandleError, mockPushToolResult, mockRemoveClosingTag)
|
||||
|
||||
expect(toolResults).toHaveLength(1)
|
||||
expect(toolResults[0]).toContain(
|
||||
`<image_data mime_type="image/webp" size="1024">${expectedBase64}</image_data>`,
|
||||
)
|
||||
})
|
||||
|
||||
it("should support SVG images", async () => {
|
||||
const imagePath = "test-image.svg"
|
||||
const imageBuffer = Buffer.from("<svg>fake-svg-data</svg>")
|
||||
const expectedBase64 = imageBuffer.toString("base64")
|
||||
|
||||
// Mock file operations
|
||||
mockIsBinaryFile.mockResolvedValue(true)
|
||||
mockFs.stat.mockResolvedValue({ size: 512 } as any)
|
||||
mockFs.readFile.mockResolvedValue(imageBuffer)
|
||||
|
||||
const block: ToolUse = {
|
||||
type: "tool_use",
|
||||
name: "read_file",
|
||||
params: {
|
||||
args: `<file><path>${imagePath}</path></file>`,
|
||||
},
|
||||
partial: false,
|
||||
}
|
||||
|
||||
await readFileTool(mockTask, block, mockAskApproval, mockHandleError, mockPushToolResult, mockRemoveClosingTag)
|
||||
|
||||
expect(toolResults).toHaveLength(1)
|
||||
expect(toolResults[0]).toContain(
|
||||
`<image_data mime_type="image/svg+xml" size="512">${expectedBase64}</image_data>`,
|
||||
)
|
||||
})
|
||||
})
|
||||
|
|
@ -1,5 +1,6 @@
|
|||
import path from "path"
|
||||
import { isBinaryFile } from "isbinaryfile"
|
||||
import fs from "fs/promises"
|
||||
|
||||
import { Task } from "../task/Task"
|
||||
import { ClineSayTool } from "../../shared/ExtensionMessage"
|
||||
|
|
@ -15,6 +16,67 @@ import { extractTextFromFile, addLineNumbers, getSupportedBinaryFormats } from "
|
|||
import { parseSourceCodeDefinitionsForFile } from "../../services/tree-sitter"
|
||||
import { parseXml } from "../../utils/xml"
|
||||
|
||||
// Supported image formats for reading as base64
|
||||
const SUPPORTED_IMAGE_FORMATS = [".png", ".jpg", ".jpeg", ".webp", ".gif", ".bmp", ".ico", ".tiff", ".svg"]
|
||||
|
||||
// Maximum file size for images (10MB) to prevent memory issues
|
||||
const MAX_IMAGE_SIZE = 10 * 1024 * 1024
|
||||
|
||||
/**
|
||||
* Checks if a file is a supported image format
|
||||
*/
|
||||
function isSupportedImageFile(filePath: string): boolean {
|
||||
const ext = path.extname(filePath).toLowerCase()
|
||||
return SUPPORTED_IMAGE_FORMATS.includes(ext)
|
||||
}
|
||||
|
||||
/**
|
||||
* Gets the MIME type for an image file
|
||||
*/
|
||||
function getImageMimeType(filePath: string): string {
|
||||
const ext = path.extname(filePath).toLowerCase()
|
||||
switch (ext) {
|
||||
case ".png":
|
||||
return "image/png"
|
||||
case ".jpg":
|
||||
case ".jpeg":
|
||||
return "image/jpeg"
|
||||
case ".webp":
|
||||
return "image/webp"
|
||||
case ".gif":
|
||||
return "image/gif"
|
||||
case ".bmp":
|
||||
return "image/bmp"
|
||||
case ".ico":
|
||||
return "image/x-icon"
|
||||
case ".tiff":
|
||||
return "image/tiff"
|
||||
case ".svg":
|
||||
return "image/svg+xml"
|
||||
default:
|
||||
return "application/octet-stream"
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Reads an image file as base64 with size validation
|
||||
*/
|
||||
async function readImageAsBase64(filePath: string): Promise<{ base64: string; mimeType: string; size: number }> {
|
||||
const stats = await fs.stat(filePath)
|
||||
|
||||
if (stats.size > MAX_IMAGE_SIZE) {
|
||||
throw new Error(
|
||||
`Image file is too large (${Math.round(stats.size / 1024 / 1024)}MB). Maximum allowed size is ${MAX_IMAGE_SIZE / 1024 / 1024}MB.`,
|
||||
)
|
||||
}
|
||||
|
||||
const buffer = await fs.readFile(filePath)
|
||||
const base64 = buffer.toString("base64")
|
||||
const mimeType = getImageMimeType(filePath)
|
||||
|
||||
return { base64, mimeType, size: stats.size }
|
||||
}
|
||||
|
||||
export function getReadFileToolDescription(blockName: string, blockParams: any): string {
|
||||
// Handle both single path and multiple files via args
|
||||
if (blockParams.args) {
|
||||
|
|
@ -68,6 +130,8 @@ interface FileResult {
|
|||
xmlContent?: string // Final XML content for this file
|
||||
feedbackText?: string // User feedback text from approval/denial
|
||||
feedbackImages?: any[] // User feedback images from approval/denial
|
||||
imageData?: string // Base64 encoded image data for supported image files
|
||||
mimeType?: string // MIME type of the image file
|
||||
}
|
||||
|
||||
export async function readFileTool(
|
||||
|
|
@ -440,6 +504,35 @@ export async function readFileTool(
|
|||
const fileExtension = path.extname(relPath).toLowerCase()
|
||||
const supportedBinaryFormats = getSupportedBinaryFormats()
|
||||
|
||||
// Check if it's a supported image file
|
||||
if (isSupportedImageFile(relPath)) {
|
||||
try {
|
||||
const { base64, mimeType, size } = await readImageAsBase64(fullPath)
|
||||
|
||||
// Track file read
|
||||
await cline.fileContextTracker.trackFileContext(relPath, "read_tool" as RecordSource)
|
||||
|
||||
updateFileResult(relPath, {
|
||||
imageData: base64,
|
||||
mimeType,
|
||||
xmlContent: `<file><path>${relPath}</path>\n<image_data mime_type="${mimeType}" size="${size}">${base64}</image_data>\n</file>`,
|
||||
})
|
||||
continue
|
||||
} catch (error) {
|
||||
const errorMsg = error instanceof Error ? error.message : String(error)
|
||||
updateFileResult(relPath, {
|
||||
status: "error",
|
||||
error: `Error reading image file: ${errorMsg}`,
|
||||
xmlContent: `<file><path>${relPath}</path><error>Error reading image file: ${errorMsg}</error></file>`,
|
||||
})
|
||||
await handleError(
|
||||
`reading image file ${relPath}`,
|
||||
error instanceof Error ? error : new Error(errorMsg),
|
||||
)
|
||||
continue
|
||||
}
|
||||
}
|
||||
|
||||
if (!supportedBinaryFormats.includes(fileExtension)) {
|
||||
updateFileResult(relPath, {
|
||||
notice: "Binary file",
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue