Skip to content

Commit 8bdd350

Browse files
committed
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
1 parent 3a8ba27 commit 8bdd350

File tree

2 files changed

+374
-0
lines changed

2 files changed

+374
-0
lines changed
Lines changed: 281 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,281 @@
1+
import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"
2+
import path from "path"
3+
import fs from "fs/promises"
4+
import { isBinaryFile } from "isbinaryfile"
5+
import { readFileTool } from "../readFileTool"
6+
import { Task } from "../../task/Task"
7+
import { ToolUse } from "../../../shared/tools"
8+
import { countFileLines } from "../../../integrations/misc/line-counter"
9+
import { extractTextFromFile, getSupportedBinaryFormats, addLineNumbers } from "../../../integrations/misc/extract-text"
10+
import { parseSourceCodeDefinitionsForFile } from "../../../services/tree-sitter"
11+
12+
// Mock dependencies
13+
vi.mock("fs/promises", () => ({
14+
default: {
15+
stat: vi.fn(),
16+
readFile: vi.fn(),
17+
},
18+
}))
19+
vi.mock("isbinaryfile")
20+
vi.mock("../../task/Task")
21+
vi.mock("../../../integrations/misc/line-counter")
22+
vi.mock("../../../integrations/misc/extract-text")
23+
vi.mock("../../../services/tree-sitter")
24+
25+
// Create mock functions
26+
const getSupportedBinaryFormatsMock = vi.fn(() => [".pdf", ".docx", ".ipynb", ".xlsx"])
27+
const extractTextFromFileMock = vi.fn()
28+
const addLineNumbersMock = vi.fn()
29+
30+
// Mock RooIgnoreController to handle vscode import
31+
vi.mock("../../ignore/RooIgnoreController", () => ({
32+
RooIgnoreController: class {
33+
initialize() {
34+
return Promise.resolve()
35+
}
36+
validateAccess() {
37+
return true
38+
}
39+
},
40+
}))
41+
42+
// Mock other dependencies
43+
vi.mock("../../../utils/fs", () => ({
44+
fileExistsAtPath: vi.fn().mockReturnValue(true),
45+
}))
46+
47+
const mockFs = vi.mocked(fs)
48+
const mockIsBinaryFile = vi.mocked(isBinaryFile)
49+
const mockedCountFileLines = vi.mocked(countFileLines)
50+
const mockedExtractTextFromFile = vi.mocked(extractTextFromFile)
51+
const mockedGetSupportedBinaryFormats = vi.mocked(getSupportedBinaryFormats)
52+
const mockedAddLineNumbers = vi.mocked(addLineNumbers)
53+
const mockedParseSourceCodeDefinitionsForFile = vi.mocked(parseSourceCodeDefinitionsForFile)
54+
55+
describe("readFileTool - Image Support", () => {
56+
let mockTask: any
57+
let mockAskApproval: any
58+
let mockHandleError: any
59+
let mockPushToolResult: any
60+
let mockRemoveClosingTag: any
61+
let toolResults: string[]
62+
63+
beforeEach(() => {
64+
vi.clearAllMocks()
65+
toolResults = []
66+
67+
// Setup mocks
68+
mockedGetSupportedBinaryFormats.mockReturnValue([".pdf", ".docx", ".ipynb", ".xlsx"])
69+
mockedCountFileLines.mockResolvedValue(0)
70+
mockedExtractTextFromFile.mockResolvedValue("")
71+
mockedAddLineNumbers.mockImplementation((text) => text)
72+
mockedParseSourceCodeDefinitionsForFile.mockResolvedValue("")
73+
74+
mockTask = {
75+
cwd: "/test/workspace",
76+
rooIgnoreController: {
77+
validateAccess: vi.fn().mockReturnValue(true),
78+
},
79+
fileContextTracker: {
80+
trackFileContext: vi.fn(),
81+
},
82+
providerRef: {
83+
deref: vi.fn().mockReturnValue({
84+
getState: vi.fn().mockResolvedValue({ maxReadFileLine: -1 }),
85+
}),
86+
},
87+
ask: vi.fn().mockResolvedValue({ response: "yesButtonClicked" }),
88+
}
89+
90+
mockAskApproval = vi.fn()
91+
mockHandleError = vi.fn()
92+
mockPushToolResult = vi.fn((result: string) => {
93+
toolResults.push(result)
94+
})
95+
mockRemoveClosingTag = vi.fn()
96+
})
97+
98+
afterEach(() => {
99+
vi.resetAllMocks()
100+
})
101+
102+
it("should read PNG image file as base64", async () => {
103+
const imagePath = "test-image.png"
104+
const imageBuffer = Buffer.from("fake-png-data")
105+
const expectedBase64 = imageBuffer.toString("base64")
106+
107+
// Mock file operations
108+
mockIsBinaryFile.mockResolvedValue(true)
109+
mockFs.stat.mockResolvedValue({ size: 1024 } as any)
110+
mockFs.readFile.mockResolvedValue(imageBuffer)
111+
112+
const block: ToolUse = {
113+
type: "tool_use",
114+
name: "read_file",
115+
params: {
116+
args: `<file><path>${imagePath}</path></file>`,
117+
},
118+
partial: false,
119+
}
120+
121+
await readFileTool(mockTask, block, mockAskApproval, mockHandleError, mockPushToolResult, mockRemoveClosingTag)
122+
123+
expect(toolResults).toHaveLength(1)
124+
expect(toolResults[0]).toContain(`<image_data mime_type="image/png" size="1024">${expectedBase64}</image_data>`)
125+
expect(mockTask.fileContextTracker.trackFileContext).toHaveBeenCalledWith(imagePath, "read_tool")
126+
})
127+
128+
it("should read JPEG image file as base64", async () => {
129+
const imagePath = "test-image.jpg"
130+
const imageBuffer = Buffer.from("fake-jpeg-data")
131+
const expectedBase64 = imageBuffer.toString("base64")
132+
133+
// Mock file operations
134+
mockIsBinaryFile.mockResolvedValue(true)
135+
mockFs.stat.mockResolvedValue({ size: 2048 } as any)
136+
mockFs.readFile.mockResolvedValue(imageBuffer)
137+
138+
const block: ToolUse = {
139+
type: "tool_use",
140+
name: "read_file",
141+
params: {
142+
args: `<file><path>${imagePath}</path></file>`,
143+
},
144+
partial: false,
145+
}
146+
147+
await readFileTool(mockTask, block, mockAskApproval, mockHandleError, mockPushToolResult, mockRemoveClosingTag)
148+
149+
expect(toolResults).toHaveLength(1)
150+
expect(toolResults[0]).toContain(
151+
`<image_data mime_type="image/jpeg" size="2048">${expectedBase64}</image_data>`,
152+
)
153+
})
154+
155+
it("should handle multiple image files", async () => {
156+
const imagePaths = ["image1.png", "image2.jpg"]
157+
const imageBuffers = [Buffer.from("fake-png-data"), Buffer.from("fake-jpeg-data")]
158+
159+
// Mock file operations
160+
mockIsBinaryFile.mockResolvedValue(true)
161+
mockFs.stat.mockResolvedValueOnce({ size: 1024 } as any).mockResolvedValueOnce({ size: 2048 } as any)
162+
mockFs.readFile.mockResolvedValueOnce(imageBuffers[0]).mockResolvedValueOnce(imageBuffers[1])
163+
164+
const block: ToolUse = {
165+
type: "tool_use",
166+
name: "read_file",
167+
params: {
168+
args: `<file><path>${imagePaths[0]}</path></file><file><path>${imagePaths[1]}</path></file>`,
169+
},
170+
partial: false,
171+
}
172+
173+
await readFileTool(mockTask, block, mockAskApproval, mockHandleError, mockPushToolResult, mockRemoveClosingTag)
174+
175+
expect(toolResults).toHaveLength(1)
176+
expect(toolResults[0]).toContain(`<image_data mime_type="image/png"`)
177+
expect(toolResults[0]).toContain(`<image_data mime_type="image/jpeg"`)
178+
expect(toolResults[0]).toContain(imageBuffers[0].toString("base64"))
179+
expect(toolResults[0]).toContain(imageBuffers[1].toString("base64"))
180+
})
181+
182+
it("should reject image files that are too large", async () => {
183+
const imagePath = "large-image.png"
184+
const largeSize = 15 * 1024 * 1024 // 15MB (exceeds 10MB limit)
185+
186+
// Mock file operations
187+
mockIsBinaryFile.mockResolvedValue(true)
188+
mockFs.stat.mockResolvedValue({ size: largeSize } as any)
189+
190+
const block: ToolUse = {
191+
type: "tool_use",
192+
name: "read_file",
193+
params: {
194+
args: `<file><path>${imagePath}</path></file>`,
195+
},
196+
partial: false,
197+
}
198+
199+
await readFileTool(mockTask, block, mockAskApproval, mockHandleError, mockPushToolResult, mockRemoveClosingTag)
200+
201+
expect(toolResults).toHaveLength(1)
202+
expect(toolResults[0]).toContain("<error>Error reading image file:")
203+
expect(toolResults[0]).toContain("Image file is too large")
204+
expect(mockHandleError).toHaveBeenCalled()
205+
})
206+
207+
it("should handle unsupported image formats as binary files", async () => {
208+
const imagePath = "test-image.tga" // Unsupported format
209+
210+
// Mock file operations
211+
mockIsBinaryFile.mockResolvedValue(true)
212+
213+
const block: ToolUse = {
214+
type: "tool_use",
215+
name: "read_file",
216+
params: {
217+
args: `<file><path>${imagePath}</path></file>`,
218+
},
219+
partial: false,
220+
}
221+
222+
await readFileTool(mockTask, block, mockAskApproval, mockHandleError, mockPushToolResult, mockRemoveClosingTag)
223+
224+
expect(toolResults).toHaveLength(1)
225+
expect(toolResults[0]).toContain("<notice>Binary file</notice>")
226+
})
227+
228+
it("should support WebP images", async () => {
229+
const imagePath = "test-image.webp"
230+
const imageBuffer = Buffer.from("fake-webp-data")
231+
const expectedBase64 = imageBuffer.toString("base64")
232+
233+
// Mock file operations
234+
mockIsBinaryFile.mockResolvedValue(true)
235+
mockFs.stat.mockResolvedValue({ size: 1024 } as any)
236+
mockFs.readFile.mockResolvedValue(imageBuffer)
237+
238+
const block: ToolUse = {
239+
type: "tool_use",
240+
name: "read_file",
241+
params: {
242+
args: `<file><path>${imagePath}</path></file>`,
243+
},
244+
partial: false,
245+
}
246+
247+
await readFileTool(mockTask, block, mockAskApproval, mockHandleError, mockPushToolResult, mockRemoveClosingTag)
248+
249+
expect(toolResults).toHaveLength(1)
250+
expect(toolResults[0]).toContain(
251+
`<image_data mime_type="image/webp" size="1024">${expectedBase64}</image_data>`,
252+
)
253+
})
254+
255+
it("should support SVG images", async () => {
256+
const imagePath = "test-image.svg"
257+
const imageBuffer = Buffer.from("<svg>fake-svg-data</svg>")
258+
const expectedBase64 = imageBuffer.toString("base64")
259+
260+
// Mock file operations
261+
mockIsBinaryFile.mockResolvedValue(true)
262+
mockFs.stat.mockResolvedValue({ size: 512 } as any)
263+
mockFs.readFile.mockResolvedValue(imageBuffer)
264+
265+
const block: ToolUse = {
266+
type: "tool_use",
267+
name: "read_file",
268+
params: {
269+
args: `<file><path>${imagePath}</path></file>`,
270+
},
271+
partial: false,
272+
}
273+
274+
await readFileTool(mockTask, block, mockAskApproval, mockHandleError, mockPushToolResult, mockRemoveClosingTag)
275+
276+
expect(toolResults).toHaveLength(1)
277+
expect(toolResults[0]).toContain(
278+
`<image_data mime_type="image/svg+xml" size="512">${expectedBase64}</image_data>`,
279+
)
280+
})
281+
})

0 commit comments

Comments
 (0)