fixed image issue

Author: liwilliam2021

src/core/webview/webviewMessageHandler.ts

@@ -137,19 +137,24 @@ export const webviewMessageHandler = async (
 	/**
 	 * Handles message editing operations with user confirmation
 	 */
-	const handleEditOperation = async (messageTs: number, editedContent: string): Promise<void> => {
+	const handleEditOperation = async (messageTs: number, editedContent: string, images?: string[]): Promise<void> => {
 		// Send message to webview to show edit confirmation dialog
 		await provider.postMessageToWebview({
 			type: "showEditMessageDialog",
 			messageTs,
 			text: editedContent,
+			images,
 		})
 	}
 
 	/**
 	 * Handles confirmed message editing from webview dialog
 	 */
-	const handleEditMessageConfirm = async (messageTs: number, editedContent: string): Promise<void> => {
+	const handleEditMessageConfirm = async (
+		messageTs: number,
+		editedContent: string,
+		images?: string[],
+	): Promise<void> => {
 		// Only proceed if we have a current cline
 		if (provider.getCurrentCline()) {
 			const currentCline = provider.getCurrentCline()!
@@ -168,6 +173,7 @@ export const webviewMessageHandler = async (
 						type: "askResponse",
 						askResponse: "messageResponse",
 						text: editedContent,
+						images,
 					})
 
 					// Don't initialize with history item for edit operations
@@ -193,11 +199,12 @@ export const webviewMessageHandler = async (
 		messageTs: number,
 		operation: "delete" | "edit",
 		editedContent?: string,
+		images?: string[],
 	): Promise<void> => {
 		if (operation === "delete") {
 			await handleDeleteOperation(messageTs)
 		} else if (operation === "edit" && editedContent) {
-			await handleEditOperation(messageTs, editedContent)
+			await handleEditOperation(messageTs, editedContent, images)
 		}
 	}
 
@@ -367,7 +374,12 @@ export const webviewMessageHandler = async (
 			break
 		case "selectImages":
 			const images = await selectImages()
-			await provider.postMessageToWebview({ type: "selectedImages", images })
+			await provider.postMessageToWebview({
+				type: "selectedImages",
+				images,
+				context: message.context,
+				messageTs: message.messageTs,
+			})
 			break
 		case "exportCurrentTask":
 			const currentTaskId = provider.getCurrentCline()?.taskId
@@ -1144,7 +1156,12 @@ export const webviewMessageHandler = async (
 				message.value &&
 				message.editedMessageContent
 			) {
-				await handleMessageModificationsOperation(message.value, "edit", message.editedMessageContent)
+				await handleMessageModificationsOperation(
+					message.value,
+					"edit",
+					message.editedMessageContent,
+					message.images,
+				)
 			}
 			break
 		}
@@ -1484,7 +1501,7 @@ export const webviewMessageHandler = async (
 			break
 		case "editMessageConfirm":
 			if (message.messageTs && message.text) {
-				await handleEditMessageConfirm(message.messageTs, message.text)
+				await handleEditMessageConfirm(message.messageTs, message.text, message.images)
 			}
 			break
 		case "getListApiConfiguration":

src/shared/ExtensionMessage.ts

@@ -160,6 +160,7 @@ export interface ExtensionMessage {
 	rulesFolderPath?: string
 	settings?: any
 	messageTs?: number
+	context?: string
 }
 
 export type ExtensionState = Pick<

src/shared/WebviewMessage.ts

@@ -199,6 +199,7 @@ export interface WebviewMessage {
 	editedMessageContent?: string
 	tab?: "settings" | "history" | "mcp" | "modes" | "chat" | "marketplace" | "account"
 	disabled?: boolean
+	context?: string
 	dataUri?: string
 	askResponse?: ClineAskResponse
 	apiConfiguration?: ProviderSettings

webview-ui/src/App.tsx

@@ -79,10 +79,12 @@ const App = () => {
 		isOpen: boolean
 		messageTs: number
 		text: string
+		images?: string[]
 	}>({
 		isOpen: false,
 		messageTs: 0,
 		text: "",
+		images: [],
 	})
 
 	const settingsRef = useRef<SettingsViewRef>(null)
@@ -145,7 +147,12 @@ const App = () => {
 			}
 
 			if (message.type === "showEditMessageDialog" && message.messageTs && message.text) {
-				setEditMessageDialogState({ isOpen: true, messageTs: message.messageTs, text: message.text })
+				setEditMessageDialogState({
+					isOpen: true,
+					messageTs: message.messageTs,
+					text: message.text,
+					images: message.images || [],
+				})
 			}
 
 			if (message.type === "acceptInput") {
@@ -253,6 +260,7 @@ const App = () => {
 						type: "editMessageConfirm",
 						messageTs: editMessageDialogState.messageTs,
 						text: editMessageDialogState.text,
+						images: editMessageDialogState.images,
 					})
 					setEditMessageDialogState((prev) => ({ ...prev, isOpen: false }))
 				}}

webview-ui/src/components/chat/ChatRow.tsx

@@ -1,4 +1,5 @@
 import React, { memo, useCallback, useEffect, useMemo, useRef, useState } from "react"
+import { appendImages } from "@src/utils/imageUtils"
 import { McpExecution } from "./McpExecution"
 import { useSize } from "react-use"
 import { useTranslation, Trans } from "react-i18next"
@@ -22,6 +23,7 @@ import { getLanguageFromPath } from "@src/utils/getLanguageFromPath"
 import { Button } from "@src/components/ui"
 
 import ChatTextArea from "./ChatTextArea"
+import { MAX_IMAGES_PER_MESSAGE } from "./ChatView"
 
 import { ToolUseBlock, ToolUseBlockHeader } from "../common/ToolUseBlock"
 import UpdateTodoListToolBlock from "./UpdateTodoListToolBlock"
@@ -122,6 +124,19 @@ export const ChatRowContent = ({
 	const [editImages, setEditImages] = useState<string[]>([])
 	const { copyWithFeedback } = useCopyToClipboard()
 
+	// Handle message events for image selection during edit mode
+	useEffect(() => {
+		const handleMessage = (event: MessageEvent) => {
+			const msg = event.data
+			if (msg.type === "selectedImages" && msg.context === "edit" && msg.messageTs === message.ts && isEditing) {
+				setEditImages((prevImages) => appendImages(prevImages, msg.images, MAX_IMAGES_PER_MESSAGE))
+			}
+		}
+
+		window.addEventListener("message", handleMessage)
+		return () => window.removeEventListener("message", handleMessage)
+	}, [isEditing, message.ts])
+
 	// Memoized callback to prevent re-renders caused by inline arrow functions
 	const handleToggleExpand = useCallback(() => {
 		onToggleExpand(message.ts)
@@ -153,13 +168,14 @@ export const ChatRowContent = ({
 			type: "submitEditedMessage",
 			value: message.ts,
 			editedMessageContent: editedContent,
+			images: editImages,
 		})
-	}, [message.ts, editedContent])
+	}, [message.ts, editedContent, editImages])
 
 	// Handle image selection for editing
 	const handleSelectImages = useCallback(() => {
-		vscode.postMessage({ type: "selectImages" })
-	}, [])
+		vscode.postMessage({ type: "selectImages", context: "edit", messageTs: message.ts })
+	}, [message.ts])
 
 	const [cost, apiReqCancelReason, apiReqStreamingFailedMessage] = useMemo(() => {
 		if (message.text !== null && message.text !== undefined && message.say === "api_req_started") {
@@ -1073,7 +1089,7 @@ export const ChatRowContent = ({
 										<Button
 											variant="ghost"
 											size="icon"
-											className="shrink-0 hidden"
+											className="shrink-0"
 											disabled={isStreaming}
 											onClick={(e) => {
 												e.stopPropagation()

webview-ui/src/components/chat/ChatView.tsx

@@ -9,6 +9,7 @@ import useSound from "use-sound"
 import { LRUCache } from "lru-cache"
 
 import { useDebounceEffect } from "@src/utils/useDebounceEffect"
+import { appendImages } from "@src/utils/imageUtils"
 
 import type { ClineAsk, ClineMessage } from "@roo-code/types"
 
@@ -712,10 +713,11 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
 					}
 					break
 				case "selectedImages":
-					const newImages = message.images ?? []
-					if (newImages.length > 0) {
+					// Only handle selectedImages if it's not for editing context
+					// When context is "edit", ChatRow will handle the images
+					if (message.context !== "edit") {
 						setSelectedImages((prevImages) =>
-							[...prevImages, ...newImages].slice(0, MAX_IMAGES_PER_MESSAGE),
+							appendImages(prevImages, message.images, MAX_IMAGES_PER_MESSAGE),
 						)
 					}
 					break

webview-ui/src/utils/imageUtils.ts

@@ -0,0 +1,17 @@
+/**
+ * Utility function to append new images to existing images array
+ * while respecting the maximum image limit
+ *
+ * @param currentImages - The current array of images
+ * @param newImages - The new images to append
+ * @param maxImages - The maximum number of images allowed
+ * @returns The updated images array
+ */
+export function appendImages(currentImages: string[], newImages: string[] | undefined, maxImages: number): string[] {
+	const imagesToAdd = newImages ?? []
+	if (imagesToAdd.length === 0) {
+		return currentImages
+	}
+
+	return [...currentImages, ...imagesToAdd].slice(0, maxImages)
+}