From 17ee88c8027707200439e2d4d225507dae7b969c Mon Sep 17 00:00:00 2001 From: daniel-lxs Date: Mon, 3 Nov 2025 14:48:09 -0500 Subject: [PATCH] fix: optimize image handling with dual-storage approach MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add imagesBase64 field to ClineMessage schema for efficient dual storage - Store both webview URIs (for display) and base64 (for API) when creating messages - Remove repeated base64→file→URI conversions on every render - Update all tools to use base64 from stored messages for API calls - Eliminates file I/O overhead during rendering and message updates - Improves performance by storing each format once instead of converting repeatedly --- packages/types/src/message.ts | 3 +- .../presentAssistantMessage.ts | 28 +++-- src/core/task/Task.ts | 73 ++++++----- src/core/tools/accessMcpResourceTool.ts | 7 +- src/core/tools/askFollowupQuestionTool.ts | 6 +- src/core/tools/attemptCompletionTool.ts | 7 +- src/core/tools/executeCommandTool.ts | 6 +- src/core/tools/readFileTool.ts | 11 +- .../misc/__tests__/imageDataUrl.spec.ts | 116 +++++++++++++++++- src/integrations/misc/imageDataUrl.ts | 86 ++++++++++++- 10 files changed, 298 insertions(+), 45 deletions(-) diff --git a/packages/types/src/message.ts b/packages/types/src/message.ts index 77c055c6e1..9cf23e386f 100644 --- a/packages/types/src/message.ts +++ b/packages/types/src/message.ts @@ -206,7 +206,8 @@ export const clineMessageSchema = z.object({ ask: clineAskSchema.optional(), say: clineSaySchema.optional(), text: z.string().optional(), - images: z.array(z.string()).optional(), + images: z.array(z.string()).optional(), // Webview URIs for frontend display + imagesBase64: z.array(z.string()).optional(), // Base64 data URLs for API calls partial: z.boolean().optional(), reasoning: z.string().optional(), conversationHistoryIndex: z.number().optional(), diff --git a/src/core/assistant-message/presentAssistantMessage.ts b/src/core/assistant-message/presentAssistantMessage.ts index 689675999f..cc8c68cf69 100644 --- a/src/core/assistant-message/presentAssistantMessage.ts +++ b/src/core/assistant-message/presentAssistantMessage.ts @@ -291,7 +291,12 @@ export async function presentAssistantMessage(cline: Task) { // Handle both messageResponse and noButtonClicked with text. if (text) { await cline.say("user_feedback", text, images) - pushToolResult(formatResponse.toolResult(formatResponse.toolDeniedWithFeedback(text), images)) + // Get base64 from the just-stored message for API call + const lastMessage = cline.clineMessages.at(-1) + const base64Images = lastMessage?.imagesBase64 + pushToolResult( + formatResponse.toolResult(formatResponse.toolDeniedWithFeedback(text), base64Images), + ) } else { pushToolResult(formatResponse.toolDenied()) } @@ -302,7 +307,12 @@ export async function presentAssistantMessage(cline: Task) { // Handle yesButtonClicked with text. if (text) { await cline.say("user_feedback", text, images) - pushToolResult(formatResponse.toolResult(formatResponse.toolApprovedWithFeedback(text), images)) + // Get base64 from the just-stored message for API call + const lastMessage = cline.clineMessages.at(-1) + const base64Images = lastMessage?.imagesBase64 + pushToolResult( + formatResponse.toolResult(formatResponse.toolApprovedWithFeedback(text), base64Images), + ) } return true @@ -396,18 +406,22 @@ export async function presentAssistantMessage(cline: Task) { ) if (response === "messageResponse") { - // Add user feedback to userContent. + // Add user feedback to chat (stores both formats) + await cline.say("user_feedback", text, images) + + // Get base64 from the just-stored message for API call + const lastMessage = cline.clineMessages.at(-1) + const base64Images = lastMessage?.imagesBase64 + + // Add user feedback to userContent for API cline.userMessageContent.push( { type: "text" as const, text: `Tool repetition limit reached. User feedback: ${text}`, }, - ...formatResponse.imageBlocks(images), + ...formatResponse.imageBlocks(base64Images), ) - // Add user feedback to chat. - await cline.say("user_feedback", text, images) - // Track tool repetition in telemetry. TelemetryService.instance.captureConsecutiveMistakeError(cline.taskId) } diff --git a/src/core/task/Task.ts b/src/core/task/Task.ts index b571c180da..290b54cfba 100644 --- a/src/core/task/Task.ts +++ b/src/core/task/Task.ts @@ -644,6 +644,9 @@ export class Task extends EventEmitter implements TaskLike { private async updateClineMessage(message: ClineMessage) { const provider = this.providerRef.deref() + + // Messages now store both formats, so no conversion needed + // The 'images' field already contains webview URIs for display await provider?.postMessageToWebview({ type: "messageUpdated", clineMessage: message }) this.emit(RooCodeEventName.Message, { action: "updated", message }) @@ -735,6 +738,7 @@ export class Task extends EventEmitter implements TaskLike { lastMessage.partial = partial lastMessage.progressStatus = progressStatus lastMessage.isProtected = isProtected + // Note: ask messages don't typically have images, so we don't update them here // TODO: Be more efficient about saving and posting only new // data or one whole message at a time so ignore partial for // saves, and only post parts of partial message instead of @@ -877,16 +881,8 @@ export class Task extends EventEmitter implements TaskLike { text: this.askResponseText, images: this.askResponseImages, } - // Normalize any image refs to base64 data URLs before handing back to callers - if (Array.isArray(result.images) && result.images.length > 0) { - try { - const { normalizeImageRefsToDataUrls } = await import("../../integrations/misc/imageDataUrl") - const normalized = await normalizeImageRefsToDataUrls(result.images) - result.images = normalized - } catch (e) { - console.error("[Task#ask] Failed to normalize image refs:", e) - } - } + // Images from askResponse are already webview URIs from the frontend, + // so no conversion needed here this.askResponse = undefined this.askResponseText = undefined this.askResponseImages = undefined @@ -1084,13 +1080,23 @@ export class Task extends EventEmitter implements TaskLike { throw new Error(`[RooCode#say] task ${this.taskId}.${this.instanceId} aborted`) } - // Ensure any image refs are normalized to base64 data URLs before persisting or sending to APIs + // Convert images to both formats for efficient dual storage + let webviewUris: string[] | undefined + let base64Images: string[] | undefined + if (Array.isArray(images) && images.length > 0) { try { const { normalizeImageRefsToDataUrls } = await import("../../integrations/misc/imageDataUrl") - images = await normalizeImageRefsToDataUrls(images) + + // Store original webview URIs/file paths for frontend + webviewUris = images + + // Convert to base64 for API calls + base64Images = await normalizeImageRefsToDataUrls(images) } catch (e) { console.error("[Task#say] Failed to normalize image refs:", e) + // Fall back to original images if conversion fails + webviewUris = images } } @@ -1104,7 +1110,8 @@ export class Task extends EventEmitter implements TaskLike { if (isUpdatingPreviousPartial) { // Existing partial message, so update it. lastMessage.text = text - lastMessage.images = images + lastMessage.images = webviewUris + lastMessage.imagesBase64 = base64Images lastMessage.partial = partial lastMessage.progressStatus = progressStatus this.updateClineMessage(lastMessage) @@ -1121,7 +1128,8 @@ export class Task extends EventEmitter implements TaskLike { type: "say", say: type, text, - images, + images: webviewUris, + imagesBase64: base64Images, partial, contextCondense, metadata: options.metadata, @@ -1137,7 +1145,8 @@ export class Task extends EventEmitter implements TaskLike { } lastMessage.text = text - lastMessage.images = images + lastMessage.images = webviewUris + lastMessage.imagesBase64 = base64Images lastMessage.partial = false lastMessage.progressStatus = progressStatus if (options.metadata) { @@ -1168,7 +1177,8 @@ export class Task extends EventEmitter implements TaskLike { type: "say", say: type, text, - images, + images: webviewUris, + imagesBase64: base64Images, contextCondense, metadata: options.metadata, }) @@ -1191,7 +1201,8 @@ export class Task extends EventEmitter implements TaskLike { type: "say", say: type, text, - images, + images: webviewUris, + imagesBase64: base64Images, checkpoint, contextCondense, }) @@ -1236,14 +1247,16 @@ export class Task extends EventEmitter implements TaskLike { await this.providerRef.deref()?.postStateToWebview() - // Convert webview URIs to base64 data URLs for backend storage (one-time conversion) - const { normalizeImageRefsToDataUrls } = await import("../../integrations/misc/imageDataUrl") - const base64Images = images ? await normalizeImageRefsToDataUrls(images) : undefined - - await this.say("text", task, base64Images) // Store base64 in backend messages + // Store the task message with both webview URIs and base64 + // This is now handled in say() method which stores both formats + await this.say("text", task, images) this.isInitialized = true - // Convert base64 to image blocks for API (no conversion needed, already base64) + // Get base64 from the stored message for API call + const lastMessage = this.clineMessages.at(-1) + const base64Images = lastMessage?.imagesBase64 + + // Convert base64 to image blocks for API const { formatResponse } = await import("../prompts/responses") let imageBlocks: Anthropic.ImageBlockParam[] = formatResponse.imageBlocks(base64Images) @@ -1509,11 +1522,11 @@ export class Task extends EventEmitter implements TaskLike { } if (responseImages && responseImages.length > 0) { - // Convert webview URIs to base64 data URLs for backend storage (one-time conversion) + // Images from user response are webview URIs, convert to base64 for API const { normalizeImageRefsToDataUrls } = await import("../../integrations/misc/imageDataUrl") const base64ResponseImages = await normalizeImageRefsToDataUrls(responseImages) - // Convert base64 to image blocks for API (no conversion needed, already base64) + // Convert base64 to image blocks for API const { formatResponse } = await import("../prompts/responses") const responseImageBlocks = formatResponse.imageBlocks(base64ResponseImages) newUserContent.push(...responseImageBlocks) @@ -1778,15 +1791,19 @@ export class Task extends EventEmitter implements TaskLike { ) if (response === "messageResponse") { + await this.say("user_feedback", text, images) + + // Get base64 from the just-stored message for API call + const lastMessage = this.clineMessages.at(-1) + const base64Images = lastMessage?.imagesBase64 + currentUserContent.push( ...[ { type: "text" as const, text: formatResponse.tooManyMistakes(text) }, - ...formatResponse.imageBlocks(images), + ...formatResponse.imageBlocks(base64Images), ], ) - await this.say("user_feedback", text, images) - // Track consecutive mistake errors in telemetry. TelemetryService.instance.captureConsecutiveMistakeError(this.taskId) } diff --git a/src/core/tools/accessMcpResourceTool.ts b/src/core/tools/accessMcpResourceTool.ts index c8a40f9236..8088d6bb34 100644 --- a/src/core/tools/accessMcpResourceTool.ts +++ b/src/core/tools/accessMcpResourceTool.ts @@ -82,7 +82,12 @@ export async function accessMcpResourceTool( }) await cline.say("mcp_server_response", resourceResultPretty, images) - pushToolResult(formatResponse.toolResult(resourceResultPretty, images)) + + // Get base64 from the just-stored message for API call + // Note: MCP images are already base64, but say() will store them in both formats + const lastMessage = cline.clineMessages.at(-1) + const base64Images = lastMessage?.imagesBase64 + pushToolResult(formatResponse.toolResult(resourceResultPretty, base64Images)) return } diff --git a/src/core/tools/askFollowupQuestionTool.ts b/src/core/tools/askFollowupQuestionTool.ts index e736936887..e3eb710768 100644 --- a/src/core/tools/askFollowupQuestionTool.ts +++ b/src/core/tools/askFollowupQuestionTool.ts @@ -78,7 +78,11 @@ export async function askFollowupQuestionTool( cline.consecutiveMistakeCount = 0 const { text, images } = await cline.ask("followup", JSON.stringify(follow_up_json), false) await cline.say("user_feedback", text ?? "", images) - pushToolResult(formatResponse.toolResult(`\n${text}\n`, images)) + + // Get base64 from the just-stored message for API call + const lastMessage = cline.clineMessages.at(-1) + const base64Images = lastMessage?.imagesBase64 + pushToolResult(formatResponse.toolResult(`\n${text}\n`, base64Images)) return } diff --git a/src/core/tools/attemptCompletionTool.ts b/src/core/tools/attemptCompletionTool.ts index 5074d7f4e8..74ca1282a8 100644 --- a/src/core/tools/attemptCompletionTool.ts +++ b/src/core/tools/attemptCompletionTool.ts @@ -121,6 +121,11 @@ export async function attemptCompletionTool( } await cline.say("user_feedback", text ?? "", images) + + // Get base64 from the just-stored message for API call + const lastMessage = cline.clineMessages.at(-1) + const base64Images = lastMessage?.imagesBase64 + const toolResults: (Anthropic.TextBlockParam | Anthropic.ImageBlockParam)[] = [] toolResults.push({ @@ -128,7 +133,7 @@ export async function attemptCompletionTool( text: `The user has provided feedback on the results. Consider their input to continue the task, and then attempt completion again.\n\n${text}\n`, }) - toolResults.push(...formatResponse.imageBlocks(images)) + toolResults.push(...formatResponse.imageBlocks(base64Images)) cline.userMessageContent.push({ type: "text", text: `${toolDescription()} Result:` }) cline.userMessageContent.push(...toolResults) diff --git a/src/core/tools/executeCommandTool.ts b/src/core/tools/executeCommandTool.ts index 2c7ce0d023..11e03a2a49 100644 --- a/src/core/tools/executeCommandTool.ts +++ b/src/core/tools/executeCommandTool.ts @@ -311,6 +311,10 @@ export async function executeCommand( const { text, images } = message await task.say("user_feedback", text, images) + // Get base64 from the just-stored message for API call + const lastMessage = task.clineMessages.at(-1) + const base64Images = lastMessage?.imagesBase64 + return [ true, formatResponse.toolResult( @@ -320,7 +324,7 @@ export async function executeCommand( `The user provided the following feedback:`, `\n${text}\n`, ].join("\n"), - images, + base64Images, ), ] } else if (completed || exitDetails) { diff --git a/src/core/tools/readFileTool.ts b/src/core/tools/readFileTool.ts index 01427f4d9d..61d445675a 100644 --- a/src/core/tools/readFileTool.ts +++ b/src/core/tools/readFileTool.ts @@ -406,10 +406,14 @@ export async function readFileTool( const { response, text, images } = await cline.ask("tool", completeMessage, false) + let feedbackBase64Images: string[] | undefined if (response !== "yesButtonClicked") { // Handle both messageResponse and noButtonClicked with text if (text) { await cline.say("user_feedback", text, images) + // Get base64 from the just-stored message + const lastMessage = cline.clineMessages.at(-1) + feedbackBase64Images = lastMessage?.imagesBase64 } cline.didRejectTool = true @@ -417,18 +421,21 @@ export async function readFileTool( status: "denied", xmlContent: `${relPath}Denied by user`, feedbackText: text, - feedbackImages: images, + feedbackImages: feedbackBase64Images, }) } else { // Handle yesButtonClicked with text if (text) { await cline.say("user_feedback", text, images) + // Get base64 from the just-stored message + const lastMessage = cline.clineMessages.at(-1) + feedbackBase64Images = lastMessage?.imagesBase64 } updateFileResult(relPath, { status: "approved", feedbackText: text, - feedbackImages: images, + feedbackImages: feedbackBase64Images, }) } } diff --git a/src/integrations/misc/__tests__/imageDataUrl.spec.ts b/src/integrations/misc/__tests__/imageDataUrl.spec.ts index bcbbe048df..7d42451115 100644 --- a/src/integrations/misc/__tests__/imageDataUrl.spec.ts +++ b/src/integrations/misc/__tests__/imageDataUrl.spec.ts @@ -1,7 +1,8 @@ -import { describe, it, expect, vi, beforeEach } from "vitest" -import { normalizeImageRefsToDataUrls } from "../imageDataUrl" +import { describe, it, expect, vi, beforeEach, afterEach } from "vitest" +import { normalizeImageRefsToDataUrls, normalizeDataUrlsToFilePaths } from "../imageDataUrl" import * as fs from "fs/promises" import * as path from "path" +import * as os from "os" // Mock fs module vi.mock("fs/promises") @@ -63,3 +64,114 @@ describe("normalizeImageRefsToDataUrls", () => { expect(result).toEqual([]) }) }) + +describe("normalizeDataUrlsToFilePaths", () => { + const testGlobalStoragePath = path.join(os.tmpdir(), "test-roo-code-storage") + + beforeEach(async () => { + vi.clearAllMocks() + // Mock mkdir to succeed + vi.mocked(fs.mkdir).mockResolvedValue(undefined) + // Mock writeFile to succeed + vi.mocked(fs.writeFile).mockResolvedValue(undefined) + // Mock access to fail initially (file doesn't exist) + vi.mocked(fs.access).mockRejectedValue(new Error("File not found")) + }) + + afterEach(() => { + vi.clearAllMocks() + }) + + it("should pass through non-data URLs unchanged", async () => { + const filePath = "/path/to/image.png" + const result = await normalizeDataUrlsToFilePaths([filePath], testGlobalStoragePath) + + expect(result).toEqual([filePath]) + expect(fs.writeFile).not.toHaveBeenCalled() + }) + + it("should convert data URLs to file paths", async () => { + const dataUrl = + "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg==" + + const result = await normalizeDataUrlsToFilePaths([dataUrl], testGlobalStoragePath) + + expect(result).toHaveLength(1) + expect(result[0]).toMatch(/temp-images/) + expect(result[0]).toMatch(/\.png$/) + expect(fs.mkdir).toHaveBeenCalledWith( + expect.stringContaining("temp-images"), + expect.objectContaining({ recursive: true }), + ) + expect(fs.writeFile).toHaveBeenCalledWith(expect.stringMatching(/temp-images.*\.png$/), expect.any(Buffer)) + }) + + it("should handle mixed arrays of data URLs and file paths", async () => { + const dataUrl = "data:image/jpeg;base64,test123" + const filePath = "/existing/image.png" + + const result = await normalizeDataUrlsToFilePaths([filePath, dataUrl], testGlobalStoragePath) + + expect(result).toHaveLength(2) + expect(result[0]).toBe(filePath) // File path unchanged + expect(result[1]).toMatch(/temp-images/) // Data URL converted + expect(result[1]).toMatch(/\.jpeg$/) // Correct extension from MIME type + }) + + it("should reuse cached file paths for same data URL", async () => { + const dataUrl = "data:image/png;base64,test123" + + // First call - should write file + const result1 = await normalizeDataUrlsToFilePaths([dataUrl], testGlobalStoragePath) + expect(fs.writeFile).toHaveBeenCalledTimes(1) + + // Mock file exists for second call + vi.mocked(fs.access).mockResolvedValue(undefined) + + // Second call - should use cached path + const result2 = await normalizeDataUrlsToFilePaths([dataUrl], testGlobalStoragePath) + + expect(result1[0]).toBe(result2[0]) // Same path returned + expect(fs.writeFile).toHaveBeenCalledTimes(1) // Not written again + }) + + it("should handle errors gracefully", async () => { + const dataUrl = "data:image/png;base64,test" + + // Mock writeFile to fail + vi.mocked(fs.writeFile).mockRejectedValue(new Error("Write failed")) + + const result = await normalizeDataUrlsToFilePaths([dataUrl], testGlobalStoragePath) + + // Should return original data URL as fallback + expect(result).toEqual([dataUrl]) + }) + + it("should handle invalid data URL format", async () => { + const invalidDataUrl = "data:image/png" // Missing base64 data + + const result = await normalizeDataUrlsToFilePaths([invalidDataUrl], testGlobalStoragePath) + + // Should skip invalid URLs + expect(result).toHaveLength(0) + }) + + it("should handle empty arrays", async () => { + const result = await normalizeDataUrlsToFilePaths([], testGlobalStoragePath) + expect(result).toEqual([]) + }) + + it("should extract correct file extensions from MIME types", async () => { + const testCases = [ + { dataUrl: "data:image/png;base64,test", expectedExt: ".png" }, + { dataUrl: "data:image/jpeg;base64,test", expectedExt: ".jpeg" }, + { dataUrl: "data:image/gif;base64,test", expectedExt: ".gif" }, + { dataUrl: "data:image/webp;base64,test", expectedExt: ".webp" }, + ] + + for (const { dataUrl, expectedExt } of testCases) { + const result = await normalizeDataUrlsToFilePaths([dataUrl], testGlobalStoragePath) + expect(result[0]).toMatch(new RegExp(`${expectedExt.replace(".", "\\.")}$`)) + } + }) +}) diff --git a/src/integrations/misc/imageDataUrl.ts b/src/integrations/misc/imageDataUrl.ts index 42d93606f7..06810aa289 100644 --- a/src/integrations/misc/imageDataUrl.ts +++ b/src/integrations/misc/imageDataUrl.ts @@ -1,6 +1,7 @@ import * as fs from "fs/promises" import * as path from "path" -import { getImageBase64ForPath } from "./image-cache" +import * as crypto from "crypto" +import { getImageBase64ForPath, setImageBase64ForPath } from "./image-cache" /** * Converts webview URIs to base64 data URLs for API calls. @@ -43,6 +44,89 @@ export async function normalizeImageRefsToDataUrls(imageRefs: string[]): Promise return results } +/** + * Cache mapping base64 data URLs to their temporary file paths + * This prevents writing the same image multiple times + */ +const base64ToFilePathCache = new Map() + +/** + * Converts base64 data URLs to file paths suitable for webview URIs. + * Writes base64 images to temporary files and returns their paths. + * This is the reverse of normalizeImageRefsToDataUrls. + */ +export async function normalizeDataUrlsToFilePaths(dataUrls: string[], globalStoragePath: string): Promise { + const results: string[] = [] + + // Ensure temp directory exists + const tempDir = path.join(globalStoragePath, "temp-images") + try { + await fs.mkdir(tempDir, { recursive: true }) + } catch (error) { + console.error("Failed to create temp-images directory:", error) + return dataUrls // Return original data URLs as fallback + } + + for (const dataUrl of dataUrls) { + // If it's not a data URL, keep it as is (might be a file path already) + if (!dataUrl.startsWith("data:image/")) { + results.push(dataUrl) + continue + } + + try { + // Check cache first + const hash = crypto.createHash("md5").update(dataUrl).digest("hex") + let filePath = base64ToFilePathCache.get(hash) + + if (filePath) { + // Check if file still exists + try { + await fs.access(filePath) + results.push(filePath) + continue + } catch { + // File was deleted, need to recreate + base64ToFilePathCache.delete(hash) + } + } + + // Extract base64 data and MIME type + const commaIndex = dataUrl.indexOf(",") + if (commaIndex === -1) { + console.error("Invalid data URL format:", dataUrl.substring(0, 50)) + continue + } + + const header = dataUrl.substring(0, commaIndex) + const base64Data = dataUrl.substring(commaIndex + 1) + + // Determine file extension from MIME type + const mimeMatch = header.match(/data:image\/([^;]+)/) + const extension = mimeMatch ? `.${mimeMatch[1]}` : ".png" + + // Write to temp file + filePath = path.join(tempDir, `${hash}${extension}`) + const buffer = Buffer.from(base64Data, "base64") + await fs.writeFile(filePath, buffer) + + // Cache the mapping + base64ToFilePathCache.set(hash, filePath) + + // Also cache in the image cache for reverse lookups + setImageBase64ForPath(filePath, dataUrl) + + results.push(filePath) + } catch (error) { + console.error("Failed to convert base64 data URL to file path:", error) + // Keep the original data URL as fallback + results.push(dataUrl) + } + } + + return results +} + /** * Converts a webview URI to a file system path */