From a6be9c4f03299f8f0eae9147b754ac52c13dc4bf Mon Sep 17 00:00:00 2001
From: Roo Code <roomote@roocode.com>
Date: Thu, 21 Aug 2025 14:55:22 +0000
Subject: [PATCH] feat: Add enhanced context editing features

- Enable Edit button for user messages
- Add 'Delete and Restore' button to deletion dialog
- Implement backend support for synchronized file state restoration
---
 src/core/webview/webviewMessageHandler.ts     | 47 +++++++++++++++++--
 src/shared/ExtensionMessage.ts                |  1 +
 src/shared/WebviewMessage.ts                  |  1 +
 webview-ui/src/App.tsx                        |  9 ++++
 webview-ui/src/components/chat/ChatRow.tsx    |  2 +-
 .../MessageModificationConfirmationDialog.tsx | 14 +++++-
 6 files changed, 68 insertions(+), 6 deletions(-)

diff --git a/src/core/webview/webviewMessageHandler.ts b/src/core/webview/webviewMessageHandler.ts
index 4dd0fee75e..68579cb168 100644
--- a/src/core/webview/webviewMessageHandler.ts
+++ b/src/core/webview/webviewMessageHandler.ts
@@ -96,21 +96,50 @@ export const webviewMessageHandler = async (
 		}
 	}
 
+	/**
+	 * Removes messages and restores file state to checkpoint
+	 */
+	const removeMessagesAndRestoreFileState = async (
+		currentCline: any,
+		messageIndex: number,
+		apiConversationHistoryIndex: number,
+		messageTs: number,
+	) => {
+		// First, delete messages
+		await removeMessagesThisAndSubsequent(currentCline, messageIndex, apiConversationHistoryIndex)
+
+		// Find the first checkpoint after this message
+		const checkpointMessage = currentCline.clineMessages.find(
+			(msg: ClineMessage) => msg.say === "checkpoint_saved" && msg.ts && msg.ts > messageTs,
+		)
+
+		if (checkpointMessage && checkpointMessage.text) {
+			// Restore to the checkpoint
+			const { checkpointRestore } = await import("../checkpoints")
+			await checkpointRestore(currentCline, {
+				ts: checkpointMessage.ts,
+				commitHash: checkpointMessage.text,
+				mode: "restore",
+			})
+		}
+	}
+
 	/**
 	 * Handles message deletion operations with user confirmation
 	 */
-	const handleDeleteOperation = async (messageTs: number): Promise<void> => {
+	const handleDeleteOperation = async (messageTs: number, withRestore: boolean = false): Promise<void> => {
 		// Send message to webview to show delete confirmation dialog
 		await provider.postMessageToWebview({
 			type: "showDeleteMessageDialog",
 			messageTs,
+			withRestore,
 		})
 	}
 
 	/**
 	 * Handles confirmed message deletion from webview dialog
 	 */
-	const handleDeleteMessageConfirm = async (messageTs: number): Promise<void> => {
+	const handleDeleteMessageConfirm = async (messageTs: number, withRestore: boolean = false): Promise<void> => {
 		// Only proceed if we have a current task.
 		if (provider.getCurrentTask()) {
 			const currentCline = provider.getCurrentTask()!
@@ -120,8 +149,18 @@ export const webviewMessageHandler = async (
 				try {
 					const { historyItem } = await provider.getTaskWithId(currentCline.taskId)
 
-					// Delete this message and all subsequent messages
-					await removeMessagesThisAndSubsequent(currentCline, messageIndex, apiConversationHistoryIndex)
+					if (withRestore) {
+						// Delete messages and restore file state
+						await removeMessagesAndRestoreFileState(
+							currentCline,
+							messageIndex,
+							apiConversationHistoryIndex,
+							messageTs,
+						)
+					} else {
+						// Delete this message and all subsequent messages
+						await removeMessagesThisAndSubsequent(currentCline, messageIndex, apiConversationHistoryIndex)
+					}
 
 					// Initialize with history item after deletion
 					await provider.createTaskWithHistoryItem(historyItem)
diff --git a/src/shared/ExtensionMessage.ts b/src/shared/ExtensionMessage.ts
index ebdc137432..43adb6ecd8 100644
--- a/src/shared/ExtensionMessage.ts
+++ b/src/shared/ExtensionMessage.ts
@@ -193,6 +193,7 @@ export interface ExtensionMessage {
 	messageTs?: number
 	context?: string
 	commands?: Command[]
+	withRestore?: boolean // For showDeleteMessageDialog
 }
 
 export type ExtensionState = Pick<
diff --git a/src/shared/WebviewMessage.ts b/src/shared/WebviewMessage.ts
index d59ccd556c..937de136de 100644
--- a/src/shared/WebviewMessage.ts
+++ b/src/shared/WebviewMessage.ts
@@ -254,6 +254,7 @@ export interface WebviewMessage {
 	visibility?: ShareVisibility // For share visibility
 	hasContent?: boolean // For checkRulesDirectoryResult
 	checkOnly?: boolean // For deleteCustomMode check
+	withRestore?: boolean // For deleteMessageConfirm with file state restore
 	codeIndexSettings?: {
 		// Global state settings
 		codebaseIndexEnabled: boolean
diff --git a/webview-ui/src/App.tsx b/webview-ui/src/App.tsx
index 3782242707..633a759a44 100644
--- a/webview-ui/src/App.tsx
+++ b/webview-ui/src/App.tsx
@@ -275,6 +275,15 @@ const App = () => {
 					vscode.postMessage({
 						type: "deleteMessageConfirm",
 						messageTs: deleteMessageDialogState.messageTs,
+						withRestore: false,
+					})
+					setDeleteMessageDialogState((prev) => ({ ...prev, isOpen: false }))
+				}}
+				onConfirmWithRestore={() => {
+					vscode.postMessage({
+						type: "deleteMessageConfirm",
+						messageTs: deleteMessageDialogState.messageTs,
+						withRestore: true,
 					})
 					setDeleteMessageDialogState((prev) => ({ ...prev, isOpen: false }))
 				}}
diff --git a/webview-ui/src/components/chat/ChatRow.tsx b/webview-ui/src/components/chat/ChatRow.tsx
index 4fa921f443..f20eec1c23 100644
--- a/webview-ui/src/components/chat/ChatRow.tsx
+++ b/webview-ui/src/components/chat/ChatRow.tsx
@@ -1090,7 +1090,7 @@ export const ChatRowContent = ({
 										<Button
 											variant="ghost"
 											size="icon"
-											className="shrink-0 hidden"
+											className="shrink-0"
 											disabled={isStreaming}
 											onClick={(e) => {
 												e.stopPropagation()
diff --git a/webview-ui/src/components/chat/MessageModificationConfirmationDialog.tsx b/webview-ui/src/components/chat/MessageModificationConfirmationDialog.tsx
index 883a34f3f3..a389f5372c 100644
--- a/webview-ui/src/components/chat/MessageModificationConfirmationDialog.tsx
+++ b/webview-ui/src/components/chat/MessageModificationConfirmationDialog.tsx
@@ -15,6 +15,7 @@ interface MessageModificationConfirmationDialogProps {
 	open: boolean
 	onOpenChange: (open: boolean) => void
 	onConfirm: () => void
+	onConfirmWithRestore?: () => void
 	type: "edit" | "delete"
 }
 
@@ -22,11 +23,13 @@ export const MessageModificationConfirmationDialog: React.FC<MessageModification
 	open,
 	onOpenChange,
 	onConfirm,
+	onConfirmWithRestore,
 	type,
 }) => {
 	const { t } = useAppTranslation()
 
 	const isEdit = type === "edit"
+	const isDelete = type === "delete"
 	const title = isEdit ? t("common:confirmation.editMessage") : t("common:confirmation.deleteMessage")
 	const description = isEdit ? t("common:confirmation.editWarning") : t("common:confirmation.deleteWarning")
 
@@ -44,8 +47,17 @@ export const MessageModificationConfirmationDialog: React.FC<MessageModification
 					<AlertDialogAction
 						onClick={onConfirm}
 						className="bg-vscode-button-background hover:bg-vscode-button-hoverBackground text-vscode-button-foreground border-vscode-button-border">
-						{t("common:confirmation.proceed")}
+						{isDelete
+							? t("common:confirmation.deleteOnly") || "Delete Messages Only"
+							: t("common:confirmation.proceed")}
 					</AlertDialogAction>
+					{isDelete && onConfirmWithRestore && (
+						<AlertDialogAction
+							onClick={onConfirmWithRestore}
+							className="bg-vscode-button-background hover:bg-vscode-button-hoverBackground text-vscode-button-foreground border-vscode-button-border">
+							{t("common:confirmation.deleteAndRestore") || "Delete Messages and Restore Files"}
+						</AlertDialogAction>
+					)}
 				</AlertDialogFooter>
 			</AlertDialogContent>
 		</AlertDialog>