Skip to content

Commit b2be7f4

Browse files
committed
feat: add support for image mentions using @ syntax
- Updated parseMentions to detect and process image file mentions - Added image processing logic with size validation and memory tracking - Modified processUserContentMentions to handle image data URLs - Updated Task.ts to pass image-related parameters to mention processing - Added comprehensive tests for image mention functionality - Fixed existing tests to work with new parseMentions return type Closes #6802
1 parent 37330b0 commit b2be7f4

File tree

5 files changed

+552
-57
lines changed

5 files changed

+552
-57
lines changed
Lines changed: 355 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,355 @@
1+
import { describe, it, expect, vi, beforeEach } from "vitest"
2+
import { parseMentions } from "../index"
3+
import { UrlContentFetcher } from "../../../services/browser/UrlContentFetcher"
4+
import * as imageHelpers from "../../tools/helpers/imageHelpers"
5+
import * as fs from "fs/promises"
6+
import path from "path"
7+
8+
// Mock the image helpers
9+
vi.mock("../../tools/helpers/imageHelpers", () => ({
10+
isSupportedImageFormat: vi.fn(),
11+
validateImageForProcessing: vi.fn(),
12+
processImageFile: vi.fn(),
13+
ImageMemoryTracker: vi.fn().mockImplementation(() => ({
14+
getTotalMemoryUsed: vi.fn().mockReturnValue(0),
15+
addMemoryUsage: vi.fn(),
16+
reset: vi.fn(),
17+
})),
18+
DEFAULT_MAX_IMAGE_FILE_SIZE_MB: 5,
19+
DEFAULT_MAX_TOTAL_IMAGE_SIZE_MB: 20,
20+
}))
21+
22+
// Mock fs
23+
vi.mock("fs/promises", () => ({
24+
default: {
25+
stat: vi.fn(),
26+
readFile: vi.fn(),
27+
},
28+
stat: vi.fn(),
29+
readFile: vi.fn(),
30+
}))
31+
32+
describe("Image Mentions", () => {
33+
let mockUrlContentFetcher: UrlContentFetcher
34+
35+
beforeEach(() => {
36+
vi.clearAllMocks()
37+
mockUrlContentFetcher = {
38+
launchBrowser: vi.fn(),
39+
closeBrowser: vi.fn(),
40+
urlToMarkdown: vi.fn(),
41+
} as any
42+
})
43+
44+
describe("parseMentions with image files", () => {
45+
it("should process image mentions and return image data URLs", async () => {
46+
const mockImageDataUrl =
47+
"data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg=="
48+
49+
// Mock image format check
50+
vi.mocked(imageHelpers.isSupportedImageFormat).mockReturnValue(true)
51+
52+
// Mock image validation
53+
vi.mocked(imageHelpers.validateImageForProcessing).mockResolvedValue({
54+
isValid: true,
55+
sizeInMB: 0.5,
56+
})
57+
58+
// Mock image processing
59+
vi.mocked(imageHelpers.processImageFile).mockResolvedValue({
60+
dataUrl: mockImageDataUrl,
61+
buffer: Buffer.from("test"),
62+
sizeInKB: 500,
63+
sizeInMB: 0.5,
64+
notice: "Image (500 KB)",
65+
})
66+
67+
// Mock file stats
68+
vi.mocked(fs.stat).mockResolvedValue({
69+
isFile: () => true,
70+
isDirectory: () => false,
71+
size: 512000,
72+
} as any)
73+
74+
const result = await parseMentions(
75+
"Check @/test/image.png for details",
76+
"/workspace",
77+
mockUrlContentFetcher,
78+
undefined,
79+
undefined,
80+
true,
81+
true,
82+
50,
83+
undefined,
84+
true, // supportsImages
85+
5, // maxImageFileSize
86+
20, // maxTotalImageSize
87+
)
88+
89+
expect(result.text).toContain("'test/image.png' (see below for image)")
90+
expect(result.text).toContain('<image_content path="test/image.png">')
91+
expect(result.text).toContain("Image (500 KB)")
92+
expect(result.images).toHaveLength(1)
93+
expect(result.images[0]).toBe(mockImageDataUrl)
94+
})
95+
96+
it("should handle multiple image mentions", async () => {
97+
const mockImageDataUrl1 = "data:image/png;base64,image1"
98+
const mockImageDataUrl2 = "data:image/jpeg;base64,image2"
99+
100+
vi.mocked(imageHelpers.isSupportedImageFormat).mockReturnValue(true)
101+
vi.mocked(imageHelpers.validateImageForProcessing).mockResolvedValue({
102+
isValid: true,
103+
sizeInMB: 0.5,
104+
})
105+
106+
vi.mocked(imageHelpers.processImageFile)
107+
.mockResolvedValueOnce({
108+
dataUrl: mockImageDataUrl1,
109+
buffer: Buffer.from("test1"),
110+
sizeInKB: 500,
111+
sizeInMB: 0.5,
112+
notice: "Image (500 KB)",
113+
})
114+
.mockResolvedValueOnce({
115+
dataUrl: mockImageDataUrl2,
116+
buffer: Buffer.from("test2"),
117+
sizeInKB: 300,
118+
sizeInMB: 0.3,
119+
notice: "Image (300 KB)",
120+
})
121+
122+
vi.mocked(fs.stat).mockResolvedValue({
123+
isFile: () => true,
124+
isDirectory: () => false,
125+
size: 512000,
126+
} as any)
127+
128+
const result = await parseMentions(
129+
"Compare @/image1.png with @/image2.jpg",
130+
"/workspace",
131+
mockUrlContentFetcher,
132+
undefined,
133+
undefined,
134+
true,
135+
true,
136+
50,
137+
undefined,
138+
true,
139+
5,
140+
20,
141+
)
142+
143+
expect(result.images).toHaveLength(2)
144+
expect(result.images[0]).toBe(mockImageDataUrl1)
145+
expect(result.images[1]).toBe(mockImageDataUrl2)
146+
expect(result.text).toContain("'image1.png' (see below for image)")
147+
expect(result.text).toContain("'image2.jpg' (see below for image)")
148+
})
149+
150+
it("should handle image size limit exceeded", async () => {
151+
vi.mocked(imageHelpers.isSupportedImageFormat).mockReturnValue(true)
152+
vi.mocked(imageHelpers.validateImageForProcessing).mockResolvedValue({
153+
isValid: false,
154+
reason: "size_limit",
155+
notice: "Image file is too large (10 MB). Maximum allowed size is 5 MB.",
156+
sizeInMB: 10,
157+
})
158+
159+
vi.mocked(fs.stat).mockResolvedValue({
160+
isFile: () => true,
161+
isDirectory: () => false,
162+
size: 10485760, // 10 MB
163+
} as any)
164+
165+
const result = await parseMentions(
166+
"Check @/large-image.png",
167+
"/workspace",
168+
mockUrlContentFetcher,
169+
undefined,
170+
undefined,
171+
true,
172+
true,
173+
50,
174+
undefined,
175+
true,
176+
5,
177+
20,
178+
)
179+
180+
expect(result.images).toHaveLength(0)
181+
expect(result.text).toContain("Image file is too large (10 MB). Maximum allowed size is 5 MB.")
182+
})
183+
184+
it("should handle model that doesn't support images", async () => {
185+
vi.mocked(imageHelpers.isSupportedImageFormat).mockReturnValue(true)
186+
vi.mocked(imageHelpers.validateImageForProcessing).mockResolvedValue({
187+
isValid: false,
188+
reason: "unsupported_model",
189+
notice: "Image file detected but current model does not support images. Skipping image processing.",
190+
})
191+
192+
vi.mocked(fs.stat).mockResolvedValue({
193+
isFile: () => true,
194+
isDirectory: () => false,
195+
size: 512000,
196+
} as any)
197+
198+
const result = await parseMentions(
199+
"Check @/image.png",
200+
"/workspace",
201+
mockUrlContentFetcher,
202+
undefined,
203+
undefined,
204+
true,
205+
true,
206+
50,
207+
undefined,
208+
false, // supportsImages = false
209+
5,
210+
20,
211+
)
212+
213+
expect(result.images).toHaveLength(0)
214+
expect(result.text).toContain("Image file detected but current model does not support images")
215+
})
216+
217+
it("should handle mixed content with images and regular files", async () => {
218+
const mockImageDataUrl = "data:image/png;base64,testimage"
219+
220+
// Mock for image file
221+
vi.mocked(imageHelpers.isSupportedImageFormat).mockImplementation((ext) => ext === ".png")
222+
223+
vi.mocked(imageHelpers.validateImageForProcessing).mockResolvedValue({
224+
isValid: true,
225+
sizeInMB: 0.5,
226+
})
227+
228+
vi.mocked(imageHelpers.processImageFile).mockResolvedValue({
229+
dataUrl: mockImageDataUrl,
230+
buffer: Buffer.from("test"),
231+
sizeInKB: 500,
232+
sizeInMB: 0.5,
233+
notice: "Image (500 KB)",
234+
})
235+
236+
// Mock file stats - need to handle both image and script file
237+
let statCallCount = 0
238+
vi.mocked(fs.stat).mockImplementation(async (path) => {
239+
statCallCount++
240+
// First call is for image.png, second is for script.js
241+
return {
242+
isFile: () => true,
243+
isDirectory: () => false,
244+
size: statCallCount === 1 ? 512000 : 100,
245+
} as any
246+
})
247+
248+
// Mock file read for text file
249+
vi.mocked(fs.readFile).mockResolvedValue("console.log('test');")
250+
251+
const result = await parseMentions(
252+
"Check @/image.png and @/script.js",
253+
"/workspace",
254+
mockUrlContentFetcher,
255+
undefined,
256+
undefined,
257+
true,
258+
true,
259+
50,
260+
undefined,
261+
true,
262+
5,
263+
20,
264+
)
265+
266+
expect(result.images).toHaveLength(1)
267+
expect(result.images[0]).toBe(mockImageDataUrl)
268+
expect(result.text).toContain("'image.png' (see below for image)")
269+
// The script.js file will have an error because we're not fully mocking the file system
270+
// but that's okay for this test - we're mainly testing that images and non-images are handled differently
271+
expect(result.text).toContain("script.js")
272+
})
273+
274+
it("should respect .rooignore for image files", async () => {
275+
vi.mocked(imageHelpers.isSupportedImageFormat).mockReturnValue(true)
276+
277+
const mockRooIgnoreController = {
278+
validateAccess: vi.fn().mockReturnValue(false),
279+
}
280+
281+
vi.mocked(fs.stat).mockResolvedValue({
282+
isFile: () => true,
283+
isDirectory: () => false,
284+
size: 512000,
285+
} as any)
286+
287+
const result = await parseMentions(
288+
"Check @/ignored-image.png",
289+
"/workspace",
290+
mockUrlContentFetcher,
291+
undefined,
292+
mockRooIgnoreController as any,
293+
true,
294+
true,
295+
50,
296+
undefined,
297+
true,
298+
5,
299+
20,
300+
)
301+
302+
expect(result.images).toHaveLength(0)
303+
expect(result.text).toContain("(Image ignored-image.png is ignored by .rooignore)")
304+
})
305+
306+
it("should handle total memory limit for multiple images", async () => {
307+
vi.mocked(imageHelpers.isSupportedImageFormat).mockReturnValue(true)
308+
309+
// First image validates successfully
310+
vi.mocked(imageHelpers.validateImageForProcessing)
311+
.mockResolvedValueOnce({
312+
isValid: true,
313+
sizeInMB: 15,
314+
})
315+
.mockResolvedValueOnce({
316+
isValid: false,
317+
reason: "memory_limit",
318+
notice: "Image skipped to avoid size limit (20MB). Current: 15MB + this file: 8MB. Try fewer or smaller images.",
319+
sizeInMB: 8,
320+
})
321+
322+
vi.mocked(imageHelpers.processImageFile).mockResolvedValue({
323+
dataUrl: "data:image/png;base64,firstimage",
324+
buffer: Buffer.from("test"),
325+
sizeInKB: 15360,
326+
sizeInMB: 15,
327+
notice: "Image (15360 KB)",
328+
})
329+
330+
vi.mocked(fs.stat).mockResolvedValue({
331+
isFile: () => true,
332+
isDirectory: () => false,
333+
size: 15728640, // 15 MB
334+
} as any)
335+
336+
const result = await parseMentions(
337+
"Check @/large1.png and @/large2.png",
338+
"/workspace",
339+
mockUrlContentFetcher,
340+
undefined,
341+
undefined,
342+
true,
343+
true,
344+
50,
345+
undefined,
346+
true,
347+
25, // maxImageFileSize
348+
20, // maxTotalImageSize
349+
)
350+
351+
expect(result.images).toHaveLength(1)
352+
expect(result.text).toContain("Image skipped to avoid size limit")
353+
})
354+
})
355+
})

0 commit comments

Comments
 (0)