diff --git a/src/core/mentions/__tests__/resolveImageMentions.spec.ts b/src/core/mentions/__tests__/resolveImageMentions.spec.ts index 9169f40ef7..c0c57d3f63 100644 --- a/src/core/mentions/__tests__/resolveImageMentions.spec.ts +++ b/src/core/mentions/__tests__/resolveImageMentions.spec.ts @@ -1,6 +1,7 @@ import * as path from "path" import { resolveImageMentions } from "../resolveImageMentions" +import { formatResponse } from "../../prompts/responses" vi.mock("../../tools/helpers/imageHelpers", () => ({ isSupportedImageFormat: vi.fn((ext: string) => @@ -11,9 +12,12 @@ vi.mock("../../tools/helpers/imageHelpers", () => ({ readImageAsDataUrlWithBuffer: vi.fn(), validateImageForProcessing: vi.fn(), ImageMemoryTracker: vi.fn().mockImplementation(function () { + let totalMemoryUsed = 0 return { - getTotalMemoryUsed: vi.fn().mockReturnValue(0), - addMemoryUsage: vi.fn(), + getTotalMemoryUsed: vi.fn(() => totalMemoryUsed), + addMemoryUsage: vi.fn((sizeInMB: number) => { + totalMemoryUsed += sizeInMB + }), } }), DEFAULT_MAX_IMAGE_FILE_SIZE_MB: 5, @@ -28,50 +32,154 @@ const mockValidateImage = vi.mocked(validateImageForProcessing) describe("resolveImageMentions", () => { beforeEach(() => { vi.clearAllMocks() + mockReadImageAsDataUrl.mockReset() + mockValidateImage.mockReset() // Default: validation passes mockValidateImage.mockResolvedValue({ isValid: true, sizeInMB: 0.1 }) }) - it("should append a data URL when a local png mention is present", async () => { - const dataUrl = `data:image/png;base64,${Buffer.from("png-bytes").toString("base64")}` - mockReadImageAsDataUrl.mockResolvedValue({ dataUrl, buffer: Buffer.from("png-bytes") }) + it.each([ + ["png", "image/png"], + ["PNG", "image/png"], + ["jpg", "image/jpeg"], + ["jpeg", "image/jpeg"], + ["gif", "image/gif"], + ["webp", "image/webp"], + ])("should append a supported %s image mention", async (extension, mimeType) => { + const buffer = Buffer.from("image-bytes") + const dataUrl = `data:${mimeType};base64,${buffer.toString("base64")}` + mockReadImageAsDataUrl.mockResolvedValue({ dataUrl, buffer }) + const text = `Please look at @/assets/cat.${extension}` const result = await resolveImageMentions({ - text: "Please look at @/assets/cat.png", + text, images: [], cwd: "/workspace", }) expect(mockValidateImage).toHaveBeenCalled() - expect(mockReadImageAsDataUrl).toHaveBeenCalledWith(path.resolve("/workspace", "assets/cat.png")) - expect(result.text).toBe("Please look at @/assets/cat.png") + expect(mockReadImageAsDataUrl).toHaveBeenCalledWith(path.resolve("/workspace", `assets/cat.${extension}`)) + expect(result.text).toBe(text) expect(result.images).toEqual([dataUrl]) + expect(formatResponse.imageBlocks(result.images)).toEqual([ + { + type: "image", + source: { type: "base64", media_type: mimeType, data: buffer.toString("base64") }, + }, + ]) }) - it("should support gif images (matching read_file)", async () => { - const dataUrl = `data:image/gif;base64,${Buffer.from("gif-bytes").toString("base64")}` - mockReadImageAsDataUrl.mockResolvedValue({ dataUrl, buffer: Buffer.from("gif-bytes") }) + it.each([ + ["svg", "image/svg+xml"], + ["bmp", "image/bmp"], + ["ico", "image/x-icon"], + ["tiff", "image/tiff"], + ["tif", "image/tiff"], + ["avif", "image/avif"], + ])("should skip %s mentions with unsupported MIME types", async (extension, mimeType) => { + const buffer = Buffer.from("unsupported-image-bytes") + const dataUrl = `data:${mimeType};base64,${buffer.toString("base64")}` + mockReadImageAsDataUrl.mockResolvedValue({ dataUrl, buffer }) + const text = `See @/image.${extension}` const result = await resolveImageMentions({ - text: "See @/animation.gif", + text, images: [], cwd: "/workspace", }) - expect(result.images).toEqual([dataUrl]) + expect(result.text).toBe(text) + expect(result.images).toEqual([]) + expect(formatResponse.imageBlocks(result.images)).toEqual([]) + }) + + it.each(["not-a-data-url", "data:;base64,aW1hZ2U="])( + "should skip image data URLs without a recognized MIME type: %s", + async (dataUrl) => { + mockReadImageAsDataUrl.mockResolvedValue({ dataUrl, buffer: Buffer.from("image-bytes") }) + + const result = await resolveImageMentions({ + text: "See @/image.png", + images: [], + cwd: "/workspace", + }) + + expect(result.images).toEqual([]) + }, + ) + + it("should preserve existing images and supported mentions when unsupported mentions are mixed in", async () => { + const buffer = Buffer.from("image-bytes") + const base64 = buffer.toString("base64") + const existingImage = `data:image/png;base64,${base64}` + const supportedImage = `data:image/webp;base64,${base64}` + mockReadImageAsDataUrl + .mockResolvedValueOnce({ dataUrl: `data:image/svg+xml;base64,${base64}`, buffer }) + .mockResolvedValueOnce({ dataUrl: supportedImage, buffer }) + .mockResolvedValueOnce({ dataUrl: `data:image/bmp;base64,${base64}`, buffer }) + const text = "Compare @/icon.svg, @/photo.webp and @/bitmap.bmp" + + const result = await resolveImageMentions({ + text, + images: [existingImage], + cwd: "/workspace", + }) + + expect(mockReadImageAsDataUrl).toHaveBeenCalledTimes(3) + expect(result.text).toBe(text) + expect(result.images).toEqual([existingImage, supportedImage]) + expect(formatResponse.imageBlocks(result.images)).toHaveLength(2) }) - it("should support svg images (matching read_file)", async () => { - const dataUrl = `data:image/svg+xml;base64,${Buffer.from("svg-bytes").toString("base64")}` - mockReadImageAsDataUrl.mockResolvedValue({ dataUrl, buffer: Buffer.from("svg-bytes") }) + it("should not count unsupported image mentions toward the total memory limit", async () => { + mockValidateImage.mockImplementation( + async (_filePath, _supportsImages, _maxImageFileSize, maxTotalImageSize, currentTotalMemoryUsed) => ({ + isValid: currentTotalMemoryUsed + 10 <= maxTotalImageSize, + sizeInMB: 10, + }), + ) + const buffer = Buffer.from("image-bytes") + const base64 = buffer.toString("base64") + const firstSupportedImage = `data:image/png;base64,${base64}` + const secondSupportedImage = `data:image/jpeg;base64,${base64}` + mockReadImageAsDataUrl + .mockResolvedValueOnce({ dataUrl: `data:image/svg+xml;base64,${base64}`, buffer }) + .mockResolvedValueOnce({ dataUrl: firstSupportedImage, buffer }) + .mockResolvedValueOnce({ dataUrl: secondSupportedImage, buffer }) const result = await resolveImageMentions({ - text: "See @/icon.svg", + text: "See @/icon.svg, @/first.png and @/second.jpg", images: [], cwd: "/workspace", + maxImageFileSize: 10, + maxTotalImageSize: 20, }) - expect(result.images).toEqual([dataUrl]) + expect(result.images).toEqual([firstSupportedImage, secondSupportedImage]) + expect(mockValidateImage).toHaveBeenNthCalledWith(2, path.resolve("/workspace", "first.png"), true, 10, 20, 0) + expect(mockValidateImage).toHaveBeenNthCalledWith(3, path.resolve("/workspace", "second.jpg"), true, 10, 20, 10) + }) + + it("should not count unsupported image mentions toward the image count limit", async () => { + const existingImages = Array.from( + { length: 19 }, + (_, index) => `data:image/png;base64,${Buffer.from(`existing-${index}`).toString("base64")}`, + ) + const buffer = Buffer.from("new-image-bytes") + const base64 = buffer.toString("base64") + const supportedImage = `data:image/png;base64,${base64}` + mockReadImageAsDataUrl + .mockResolvedValueOnce({ dataUrl: `data:image/svg+xml;base64,${base64}`, buffer }) + .mockResolvedValueOnce({ dataUrl: supportedImage, buffer }) + + const result = await resolveImageMentions({ + text: "See @/icon.svg, @/last.png and @/over-limit.jpg", + images: existingImages, + cwd: "/workspace", + }) + + expect(result.images).toEqual([...existingImages, supportedImage]) + expect(mockReadImageAsDataUrl).toHaveBeenCalledTimes(2) }) it("should ignore non-image mentions", async () => { diff --git a/src/core/mentions/resolveImageMentions.ts b/src/core/mentions/resolveImageMentions.ts index 0a0344348f..b4b905c5c2 100644 --- a/src/core/mentions/resolveImageMentions.ts +++ b/src/core/mentions/resolveImageMentions.ts @@ -1,6 +1,7 @@ import * as path from "path" import { mentionRegexGlobal, unescapeSpaces } from "../../shared/context-mentions" +import { getImageMimeType, isSupportedImageMimeType } from "../../utils/imageMime" import { isSupportedImageFormat, readImageAsDataUrlWithBuffer, @@ -50,8 +51,11 @@ function dedupePreserveOrder(values: string[]): string[] { * Resolves local image file mentions like `@/path/to/image.png` found in `text` into `data:image/...;base64,...` * and appends them to the outgoing `images` array. * - * Behavior matches the read_file tool: - * - Supports the same image formats: png, jpg, jpeg, gif, webp, svg, bmp, ico, tiff, avif + * Recognizes the same image extensions as the read_file tool, but only appends images + * supported by the shared image-block contract (JPEG, PNG, GIF, and WebP). + * Other image MIME types are skipped without changing the original text. + * + * Like the read_file tool: * - Respects per-file size limits (default 5MB) * - Respects total memory limits (default 20MB) * - Skips images if model doesn't support them @@ -128,6 +132,10 @@ export async function resolveImageMentions({ } const { dataUrl } = await readImageAsDataUrlWithBuffer(absPath) + if (!isSupportedImageMimeType(getImageMimeType(dataUrl))) { + // Fail-soft: skip formats unsupported by the shared image-block contract. + continue + } newImages.push(dataUrl) // Track memory usage diff --git a/src/core/prompts/__tests__/responses-images.spec.ts b/src/core/prompts/__tests__/responses-images.spec.ts new file mode 100644 index 0000000000..c88d76acf7 --- /dev/null +++ b/src/core/prompts/__tests__/responses-images.spec.ts @@ -0,0 +1,122 @@ +import { formatResponse } from "../responses" +import { IMAGE_MIME_TYPES } from "../../tools/helpers/imageHelpers" +import { getImageMimeType, isSupportedImageMimeType } from "../../../utils/imageMime" + +const supportedMimeTypes = ["image/jpeg", "image/png", "image/gif", "image/webp"] as const +const supportedMimeTypeSet = new Set(supportedMimeTypes) +const unsupportedMimeTypes = [ + ...new Set([...Object.values(IMAGE_MIME_TYPES), "image/heic", "application/octet-stream", "image/svg+xml"]), +].filter((mimeType) => !supportedMimeTypeSet.has(mimeType)) +const payload = "c3ludGhldGlj" + +describe("image MIME utilities", () => { + it.each(supportedMimeTypes)("extracts and accepts the SDK MIME type %s", (mimeType) => { + expect(getImageMimeType(`data:${mimeType};base64,${payload}`)).toBe(mimeType) + expect(isSupportedImageMimeType(mimeType)).toBe(true) + }) + + it.each(unsupportedMimeTypes)("extracts but rejects the unsupported MIME type %s", (mimeType) => { + expect(getImageMimeType(`data:${mimeType};base64,${payload}`)).toBe(mimeType) + expect(isSupportedImageMimeType(mimeType)).toBe(false) + }) + + it("returns no MIME for a missing header and rejects absent or empty MIME values", () => { + expect(getImageMimeType("")).toBeUndefined() + expect(isSupportedImageMimeType(undefined)).toBe(false) + expect(isSupportedImageMimeType("")).toBe(false) + }) +}) + +describe("shared image-block MIME contract", () => { + it.each(supportedMimeTypes)("preserves the complete image block for %s through both entry points", (mimeType) => { + const dataUrl = `data:${mimeType};base64,${payload}` + const image = { type: "image", source: { type: "base64", media_type: mimeType, data: payload } } + expect(formatResponse.imageBlocks([dataUrl])).toEqual([image]) + expect(formatResponse.toolResult("read result", [dataUrl])).toEqual([ + { type: "text", text: "read result" }, + image, + ]) + }) + + it.each(unsupportedMimeTypes)("rejects %s through both entry points before returning a payload", (mimeType) => { + const dataUrl = `data:${mimeType};base64,${payload}` + const error = new Error(`Unsupported image MIME type: ${mimeType}. Use JPEG, PNG, GIF, or WebP.`) + expect(() => formatResponse.imageBlocks([dataUrl])).toThrowError(error) + expect(() => formatResponse.toolResult("read result", [dataUrl])).toThrowError(error) + }) + + it.each([ + "", + "not a data URL", + "data:;base64,c3ludGhldGlj", + "data:image/png,c3ludGhldGlj", + "data:image/png;base64", + "data:image/png;base64;c3ludGhldGlj", + "data:image/png;charset=utf-8;base64,c3ludGhldGlj", + "DATA:image/png;base64,c3ludGhldGlj", + "data:image/png;BASE64,c3ludGhldGlj", + "prefix data:image/png;base64,c3ludGhldGlj", + " data:image/png;base64,c3ludGhldGlj", + "\uFEFFdata:image/png;base64,c3ludGhldGlj", + "data\uFF1Aimage/png;base64,c3ludGhldGlj", + "data:image/png;base64\uFF0Cc3ludGhldGlj", + ])("rejects a malformed or noncanonical header: %j", (dataUrl) => { + const error = new Error("Unsupported image MIME type: unknown. Use JPEG, PNG, GIF, or WebP.") + expect(() => formatResponse.imageBlocks([dataUrl])).toThrowError(error) + expect(() => formatResponse.toolResult("read result", [dataUrl])).toThrowError(error) + }) + + it.each(["image/PNG", "image/jpg", "image/png ", "image/p\u200Bng", "image/\u0440ng", "image\uFF0Fpng"])( + "does not normalize unsupported MIME spelling %j into an allowed type", + (mimeType) => { + const dataUrl = `data:${mimeType};base64,${payload}` + const error = new Error(`Unsupported image MIME type: ${mimeType}. Use JPEG, PNG, GIF, or WebP.`) + expect(() => formatResponse.imageBlocks([dataUrl])).toThrowError(error) + expect(() => formatResponse.toolResult("read result", [dataUrl])).toThrowError(error) + }, + ) + + it.each([{ images: undefined }, { images: [] }])("preserves the no-image response for $images", ({ images }) => { + expect(formatResponse.imageBlocks(images)).toEqual([]) + expect(formatResponse.toolResult("read result", images)).toBe("read result") + expect(formatResponse.toolResult("", images)).toBe("") + }) + + it("preserves image order, duplicate images, Unicode feedback, and the input array", () => { + const images = [ + "data:image/png;base64,Zmlyc3Q=", + "data:image/jpeg;base64,c2Vjb25k", + "data:image/png;base64,Zmlyc3Q=", + ] + const originalImages = [...images] + const expected = [ + { type: "image", source: { type: "base64", media_type: "image/png", data: "Zmlyc3Q=" } }, + { type: "image", source: { type: "base64", media_type: "image/jpeg", data: "c2Vjb25k" } }, + { type: "image", source: { type: "base64", media_type: "image/png", data: "Zmlyc3Q=" } }, + ] + const text = "\u041E\u0442\u0432\u0435\u0442 \uD83D\uDDBC\uFE0F e\u0301" + expect(formatResponse.imageBlocks(images)).toEqual(expected) + expect(formatResponse.toolResult(text, images)).toEqual([{ type: "text", text }, ...expected]) + expect(images).toEqual(originalImages) + }) + + it.each([0, 1, 2])("rejects an unsupported image at index %i without changing the input", (index) => { + const images = Array(3).fill(`data:image/png;base64,${payload}`) + images[index] = `data:image/bmp;base64,${payload}` + const originalImages = [...images] + const error = new Error("Unsupported image MIME type: image/bmp. Use JPEG, PNG, GIF, or WebP.") + expect(() => formatResponse.imageBlocks(images)).toThrowError(error) + expect(() => formatResponse.toolResult("read result", images)).toThrowError(error) + expect(images).toEqual(originalImages) + }) + + it.each(["", "not decoded here", "first,second", "\uD83D\uDE00e\u0301"])( + "preserves the opaque payload %j without adding byte or base64 validation", + (data) => { + const dataUrl = `data:image/png;base64,${data}` + const image = { type: "image", source: { type: "base64", media_type: "image/png", data } } + expect(formatResponse.imageBlocks([dataUrl])).toEqual([image]) + expect(formatResponse.toolResult("", [dataUrl])).toEqual([{ type: "text", text: "" }, image]) + }, + ) +}) diff --git a/src/core/prompts/responses.ts b/src/core/prompts/responses.ts index 7714438431..d8a3fa4fee 100644 --- a/src/core/prompts/responses.ts +++ b/src/core/prompts/responses.ts @@ -3,6 +3,7 @@ import * as path from "path" import * as diff from "diff" import { RooIgnoreController, LOCK_TEXT_SYMBOL } from "../ignore/RooIgnoreController" import { RooProtectedController } from "../protect/RooProtectedController" +import { getImageMimeType, isSupportedImageMimeType } from "../../utils/imageMime" export const formatResponse = { toolDenied: () => @@ -226,14 +227,19 @@ Otherwise, if you have not completed the task and do not need additional informa // to avoid circular dependency const formatImagesIntoBlocks = (images?: string[]): Anthropic.ImageBlockParam[] => { return images - ? images.map((dataUrl) => { + ? images.map((dataUrl): Anthropic.ImageBlockParam => { // data:image/png;base64,base64string - const [rest, base64] = dataUrl.split(",") - const mimeType = rest.split(":")[1].split(";")[0] + const mimeType = getImageMimeType(dataUrl) + if (!isSupportedImageMimeType(mimeType)) { + throw new Error( + `Unsupported image MIME type: ${mimeType ?? "unknown"}. Use JPEG, PNG, GIF, or WebP.`, + ) + } + const base64 = dataUrl.slice(dataUrl.indexOf(",") + 1) return { type: "image", source: { type: "base64", media_type: mimeType, data: base64 }, - } as Anthropic.ImageBlockParam + } }) : [] } diff --git a/src/utils/imageMime.ts b/src/utils/imageMime.ts new file mode 100644 index 0000000000..d9f5d684e6 --- /dev/null +++ b/src/utils/imageMime.ts @@ -0,0 +1,18 @@ +import type { Anthropic } from "@anthropic-ai/sdk" + +export function getImageMimeType(dataUrl: string): string | undefined { + return /^data:([^;]+);base64,/.exec(dataUrl)?.[1] +} + +/** + * All providers consume the shared Anthropic image-block contract. Since model + * capabilities expose only supportsImages, conservatively use its four MIME + * types rather than assuming provider-specific support for other encodings. + */ +export function isSupportedImageMimeType( + mimeType: string | undefined, +): mimeType is Anthropic.Base64ImageSource["media_type"] { + return ( + mimeType === "image/jpeg" || mimeType === "image/png" || mimeType === "image/gif" || mimeType === "image/webp" + ) +}