temp push

src/core/task/Task.ts

@@ -177,6 +177,7 @@ export class Task extends EventEmitter<ClineEvents> {
 	// LLM Messages & Chat Messages
 	apiConversationHistory: ApiMessage[] = []
 	clineMessages: ClineMessage[] = []
+	public pendingUserMessageCheckpoint?: Record<string, unknown>
 
 	// Ask
 	private askResponse?: ClineAskResponse
@@ -351,11 +352,30 @@ export class Task extends EventEmitter<ClineEvents> {
 	// Cline Messages
 
 	private async getSavedClineMessages(): Promise<ClineMessage[]> {
-		return readTaskMessages({ taskId: this.taskId, globalStoragePath: this.globalStoragePath })
+		const messages = await readTaskMessages({ taskId: this.taskId, globalStoragePath: this.globalStoragePath })
+		console.log("[Task#getSavedClineMessages] Loaded messages from disk:", messages.length)
+		const messagesWithCheckpoints = messages.filter((m) => m.checkpoint)
+		console.log("[Task#getSavedClineMessages] Messages with checkpoints:", messagesWithCheckpoints.length)
+		if (messagesWithCheckpoints.length > 0) {
+			console.log("[Task#getSavedClineMessages] Sample checkpoint:", messagesWithCheckpoints[0].checkpoint)
+		}
+		return messages
 	}
 
 	private async addToClineMessages(message: ClineMessage) {
+		console.log("[Task#addToClineMessages] Adding message:", JSON.stringify(message, null, 2))
 		this.clineMessages.push(message)
+
+		// Verify the message was added correctly
+		const addedMessage = this.clineMessages[this.clineMessages.length - 1]
+		console.log("[Task#addToClineMessages] Verified added message:", {
+			ts: addedMessage.ts,
+			type: addedMessage.type,
+			say: addedMessage.say,
+			hasCheckpoint: !!addedMessage.checkpoint,
+			checkpoint: addedMessage.checkpoint,
+		})
+
 		const provider = this.providerRef.deref()
 		await provider?.postStateToWebview()
 		this.emit("message", { action: "created", message })
@@ -532,6 +552,39 @@ export class Task extends EventEmitter<ClineEvents> {
 	}
 
 	async handleWebviewAskResponse(askResponse: ClineAskResponse, text?: string, images?: string[]) {
+		// Save checkpoint BEFORE setting the response to ensure it's ready when the user_feedback message is created
+		if (this.enableCheckpoints && askResponse === "messageResponse") {
+			console.log("[Task#handleWebviewAskResponse] Saving checkpoint for user message")
+			try {
+				const checkpointResult = await this.checkpointSave(true) // Force checkpoint save
+				console.log("[Task#handleWebviewAskResponse] Checkpoint result:", checkpointResult)
+				if (checkpointResult?.commit) {
+					// Store checkpoint data temporarily to be used when creating the user_feedback message
+					this.pendingUserMessageCheckpoint = {
+						hash: checkpointResult.commit,
+						timestamp: Date.now(),
+						type: "user_message",
+					}
+					console.log(
+						"[Task#handleWebviewAskResponse] Set pendingUserMessageCheckpoint:",
+						this.pendingUserMessageCheckpoint,
+					)
+				} else {
+					console.log("[Task#handleWebviewAskResponse] No commit in checkpoint result")
+				}
+			} catch (error) {
+				console.error("[Task#handleWebviewAskResponse] Failed to save checkpoint after user message:", error)
+			}
+		} else {
+			console.log(
+				"[Task#handleWebviewAskResponse] Skipping checkpoint save - enableCheckpoints:",
+				this.enableCheckpoints,
+				"askResponse:",
+				askResponse,
+			)
+		}
+
+		// Now set the response, which will trigger the ask promise to resolve
 		this.askResponse = askResponse
 		this.askResponseText = text
 		this.askResponseImages = images
@@ -705,15 +758,49 @@ export class Task extends EventEmitter<ClineEvents> {
 				this.lastMessageTs = sayTs
 			}
 
-			await this.addToClineMessages({
+			if (type === "user_feedback") {
-				ts: sayTs,
+				// Automatically use and clear the pending checkpoint for user_feedback messages
-				type: "say",
+				const feedbackCheckpoint = checkpoint || this.pendingUserMessageCheckpoint
-				say: type,
+				this.pendingUserMessageCheckpoint = undefined // Clear it after use
-				text,
+
-				images,
+				console.log("[Task#say] Adding user_feedback message with checkpoint:", feedbackCheckpoint)
-				checkpoint,
+				console.log(
-				contextCondense,
+					"[Task#say] Full message object:",
-			})
+					JSON.stringify(
+						{
+							ts: sayTs,
+							type: "say",
+							say: type,
+							text,
+							images,
+							checkpoint: feedbackCheckpoint,
+							contextCondense,
+						},
+						null,
+						2,
+					),
+				)
+
+				await this.addToClineMessages({
+					ts: sayTs,
+					type: "say",
+					say: type,
+					text,
+					images,
+					checkpoint: feedbackCheckpoint,
+					contextCondense,
+				})
+			} else {
+				await this.addToClineMessages({
+					ts: sayTs,
+					type: "say",
+					say: type,
+					text,
+					images,
+					checkpoint,
+					contextCondense,
+				})
+			}
 		}
 	}
 
@@ -740,6 +827,7 @@ export class Task extends EventEmitter<ClineEvents> {
 		this.apiConversationHistory = []
 		await this.providerRef.deref()?.postStateToWebview()
 
+		// Checkpoint will be saved in handleWebviewAskResponse before this message is created
 		await this.say("text", task, images)
 		this.isInitialized = true
 
@@ -783,6 +871,18 @@ export class Task extends EventEmitter<ClineEvents> {
 	private async resumeTaskFromHistory() {
 		const modifiedClineMessages = await this.getSavedClineMessages()
 
+		// Debug: Check if any messages have checkpoints
+		const messagesWithCheckpoints = modifiedClineMessages.filter((m) => m.checkpoint)
+		console.log("[Task#resumeTaskFromHistory] Total messages loaded:", modifiedClineMessages.length)
+		console.log("[Task#resumeTaskFromHistory] Messages with checkpoints:", messagesWithCheckpoints.length)
+		messagesWithCheckpoints.forEach((msg, idx) => {
+			console.log(`[Task#resumeTaskFromHistory] Message ${idx} with checkpoint:`, {
+				ts: msg.ts,
+				say: msg.say,
+				checkpoint: msg.checkpoint,
+			})
+		})
+
 		// Remove any resume messages that may have been added before
 		const lastRelevantMessageIndex = findLastIndex(
 			modifiedClineMessages,
@@ -836,11 +936,23 @@ export class Task extends EventEmitter<ClineEvents> {
 		let responseText: string | undefined
 		let responseImages: string[] | undefined
 		if (response === "messageResponse") {
+			// The checkpoint was already saved in handleWebviewAskResponse and attached to pendingUserMessageCheckpoint
+			// The say method will automatically handle it for user_feedback messages
+			console.log("[Task#resumeTaskFromHistory] Adding user_feedback message")
 			await this.say("user_feedback", text, images)
+
+			// Verify the message was added with checkpoint
+			const lastMessage = this.clineMessages[this.clineMessages.length - 1]
+			console.log("[Task#resumeTaskFromHistory] Last message after say:", {
+				ts: lastMessage?.ts,
+				say: lastMessage?.say,
+				hasCheckpoint: !!lastMessage?.checkpoint,
+				checkpoint: lastMessage?.checkpoint,
+			})
+
 			responseText = text
 			responseImages = images
 		}
-
 		// Make sure that the api conversation history can be resumed by the API,
 		// even if it goes out of sync with cline messages.
 		let existingApiConversationHistory: ApiMessage[] = await this.getSavedApiConversationHistory()
@@ -1171,8 +1283,19 @@ export class Task extends EventEmitter<ClineEvents> {
 					],
 				)
 
+				// The say method will automatically handle the pending checkpoint for user_feedback messages
+				console.log("[Task] Tool approval - About to say user_feedback")
 				await this.say("user_feedback", text, images)
 
+				// Verify the message was added with checkpoint
+				const lastMessage = this.clineMessages[this.clineMessages.length - 1]
+				console.log("[Task] Tool approval - Last message after say:", {
+					ts: lastMessage?.ts,
+					say: lastMessage?.say,
+					hasCheckpoint: !!lastMessage?.checkpoint,
+					checkpoint: lastMessage?.checkpoint,
+				})
+
 				// Track consecutive mistake errors in telemetry.
 				TelemetryService.instance.captureConsecutiveMistakeError(this.taskId)
 			}

src/core/webview/ClineProvider.ts

@@ -569,6 +569,10 @@ export class ClineProvider
 			`[subtasks] ${cline.parentTask ? "child" : "parent"} task ${cline.taskId}.${cline.instanceId} instantiated`,
 		)
 
+		// Don't save checkpoint here - it will be saved in handleWebviewAskResponse
+		// when the user message is actually sent. The checkpoint service might not
+		// be initialized yet at this point.
+
 		return cline
 	}
 
@@ -1442,7 +1446,23 @@ export class ClineProvider
 			currentTaskItem: this.getCurrentCline()?.taskId
 				? (taskHistory || []).find((item: HistoryItem) => item.id === this.getCurrentCline()?.taskId)
 				: undefined,
-			clineMessages: this.getCurrentCline()?.clineMessages || [],
+			clineMessages: (() => {
+				const messages = this.getCurrentCline()?.clineMessages || []
+				const messagesWithCheckpoints = messages.filter((m) => m.checkpoint)
+				console.log("[ClineProvider#getStateToPostToWebview] Total messages:", messages.length)
+				console.log(
+					"[ClineProvider#getStateToPostToWebview] Messages with checkpoints:",
+					messagesWithCheckpoints.length,
+				)
+				if (messagesWithCheckpoints.length > 0) {
+					console.log("[ClineProvider#getStateToPostToWebview] Sample message with checkpoint:", {
+						ts: messagesWithCheckpoints[0].ts,
+						say: messagesWithCheckpoints[0].say,
+						checkpoint: messagesWithCheckpoints[0].checkpoint,
+					})
+				}
+				return messages
+			})(),
 			taskHistory: (taskHistory || [])
 				.filter((item: HistoryItem) => item.ts && item.task)
 				.sort((a: HistoryItem, b: HistoryItem) => b.ts - a.ts),

src/core/webview/__tests__/webviewMessageHandler.spec.ts

@@ -565,7 +565,10 @@ describe("webviewMessageHandler - message dialog preferences", () => {
 	describe("deleteMessage", () => {
 		it("should show dialog when skipDeleteMessageConfirmation is false", async () => {
 			vi.mocked(mockClineProvider.contextProxy.getValue).mockReturnValue(false)
-			vi.mocked(mockClineProvider.getCurrentCline).mockReturnValue({} as any) // Mock current cline exists
+			vi.mocked(mockClineProvider.getCurrentCline).mockReturnValue({
+				clineMessages: [],
+				apiConversationHistory: [],
+			} as any) // Mock current cline with empty arrays
 
 			await webviewMessageHandler(mockClineProvider, {
 				type: "deleteMessage",
@@ -607,7 +610,10 @@ describe("webviewMessageHandler - message dialog preferences", () => {
 	describe("submitEditedMessage", () => {
 		it("should show dialog when skipEditMessageConfirmation is false", async () => {
 			vi.mocked(mockClineProvider.contextProxy.getValue).mockReturnValue(false)
-			vi.mocked(mockClineProvider.getCurrentCline).mockReturnValue({} as any) // Mock current cline exists
+			vi.mocked(mockClineProvider.getCurrentCline).mockReturnValue({
+				clineMessages: [],
+				apiConversationHistory: [],
+			} as any) // Mock current cline with empty arrays
 
 			await webviewMessageHandler(mockClineProvider, {
 				type: "submitEditedMessage",

src/core/webview/webviewMessageHandler.ts

@@ -106,10 +106,46 @@ export const webviewMessageHandler = async (
 			// Directly handle the deletion without showing dialog
 			await handleDeleteMessageConfirm(messageTs)
 		} else {
+			// Check if the message has a checkpoint
+			const currentCline = provider.getCurrentCline()
+			let hasCheckpoint = false
+			if (currentCline) {
+				// Debug: Log all messages to understand the state
+				console.log("[webviewMessageHandler] Total messages:", currentCline.clineMessages.length)
+				console.log("[webviewMessageHandler] Looking for message with ts:", messageTs)
+				console.log(
+					"[webviewMessageHandler] All messages with timestamps:",
+					currentCline.clineMessages.map((m, idx) => ({
+						index: idx,
+						ts: m.ts,
+						say: m.say,
+						hasCheckpoint: !!m.checkpoint,
+						checkpoint: m.checkpoint,
+					})),
+				)
+
+				const { messageIndex } = findMessageIndices(messageTs, currentCline)
+				console.log("[webviewMessageHandler] Checking for checkpoint at messageIndex:", messageIndex)
+				if (messageIndex !== -1) {
+					const targetMessage = currentCline.clineMessages[messageIndex]
+					console.log("[webviewMessageHandler] Target message:", JSON.stringify(targetMessage, null, 2))
+					console.log("[webviewMessageHandler] Target message checkpoint:", targetMessage?.checkpoint)
+					hasCheckpoint = !!(
+						targetMessage?.checkpoint &&
+						typeof targetMessage.checkpoint === "object" &&
+						"hash" in targetMessage.checkpoint
+					)
+					console.log("[webviewMessageHandler] hasCheckpoint:", hasCheckpoint)
+				} else {
+					console.log("[webviewMessageHandler] Message not found! Looking for ts:", messageTs)
+				}
+			}
+
 			// Send message to webview to show delete confirmation dialog
 			await provider.postMessageToWebview({
 				type: "showDeleteMessageDialog",
 				messageTs,
+				hasCheckpoint,
 			})
 		}
 	}
@@ -117,7 +153,7 @@ export const webviewMessageHandler = async (
 	/**
 	 * Handles confirmed message deletion from webview dialog
 	 */
-	const handleDeleteMessageConfirm = async (messageTs: number): Promise<void> => {
+	const handleDeleteMessageConfirm = async (messageTs: number, restoreCheckpoint?: boolean): Promise<void> => {
 		// Only proceed if we have a current cline
 		if (provider.getCurrentCline()) {
 			const currentCline = provider.getCurrentCline()!
@@ -125,6 +161,22 @@ export const webviewMessageHandler = async (
 
 			if (messageIndex !== -1) {
 				try {
+					// If checkpoint restoration is requested, restore to the checkpoint first
+					if (restoreCheckpoint) {
+						const targetMessage = currentCline.clineMessages[messageIndex]
+						if (
+							targetMessage?.checkpoint &&
+							typeof targetMessage.checkpoint === "object" &&
+							"hash" in targetMessage.checkpoint
+						) {
+							await currentCline.checkpointRestore({
+								ts: targetMessage.ts!,
+								commitHash: targetMessage.checkpoint.hash as string,
+								mode: "restore",
+							})
+						}
+					}
+
 					const { historyItem } = await provider.getTaskWithId(currentCline.taskId)
 
 					// Delete this message and all subsequent messages
@@ -149,15 +201,62 @@ export const webviewMessageHandler = async (
 		// Check if user has opted to skip the confirmation
 		const skipEditMessageConfirmation = getGlobalState("skipEditMessageConfirmation")
 
+		// Always check if the message has a checkpoint first
+		const currentCline = provider.getCurrentCline()
+		let hasCheckpoint = false
+		if (currentCline) {
+			console.log(
+				"[webviewMessageHandler] Edit - Total messages in currentCline:",
+				currentCline.clineMessages.length,
+			)
+			console.log("[webviewMessageHandler] Edit - Looking for messageTs:", messageTs)
+
+			// Log all messages with their timestamps and checkpoint status
+			currentCline.clineMessages.forEach((msg, idx) => {
+				console.log(
+					`[webviewMessageHandler] Edit - Message ${idx}: ts=${msg.ts}, type=${msg.type}, say=${msg.say}, hasCheckpoint=${!!msg.checkpoint}, checkpoint=${JSON.stringify(msg.checkpoint)}`,
+				)
+			})
+
+			const { messageIndex } = findMessageIndices(messageTs, currentCline)
+			console.log("[webviewMessageHandler] Edit - Checking for checkpoint at messageIndex:", messageIndex)
+			if (messageIndex !== -1) {
+				const targetMessage = currentCline.clineMessages[messageIndex]
+				console.log("[webviewMessageHandler] Edit - Target message:", JSON.stringify(targetMessage, null, 2))
+				console.log("[webviewMessageHandler] Edit - Target message checkpoint:", targetMessage?.checkpoint)
+				hasCheckpoint = !!(
+					targetMessage?.checkpoint &&
+					typeof targetMessage.checkpoint === "object" &&
+					"hash" in targetMessage.checkpoint
+				)
+				console.log("[webviewMessageHandler] Edit - hasCheckpoint:", hasCheckpoint)
+			} else {
+				console.log("[webviewMessageHandler] Edit - Message not found in clineMessages!")
+			}
+		} else {
+			console.log("[webviewMessageHandler] Edit - No currentCline available!")
+		}
+
 		if (skipEditMessageConfirmation) {
-			// Directly handle the edit without showing dialog
+			// If there's a checkpoint, show the checkpoint dialog even when skipping confirmation
-			await handleEditMessageConfirm(messageTs, editedContent)
+			if (hasCheckpoint) {
+				await provider.postMessageToWebview({
+					type: "showEditMessageDialog",
+					messageTs,
+					text: editedContent,
+					hasCheckpoint,
+				})
+			} else {
+				// No checkpoint, directly handle the edit without showing dialog
+				await handleEditMessageConfirm(messageTs, editedContent, false)
+			}
 		} else {
 			// Send message to webview to show edit confirmation dialog
 			await provider.postMessageToWebview({
 				type: "showEditMessageDialog",
 				messageTs,
 				text: editedContent,
+				hasCheckpoint,
 			})
 		}
 	}
@@ -165,7 +264,11 @@ export const webviewMessageHandler = async (
 	/**
 	 * Handles confirmed message editing from webview dialog
 	 */
-	const handleEditMessageConfirm = async (messageTs: number, editedContent: string): Promise<void> => {
+	const handleEditMessageConfirm = async (
+		messageTs: number,
+		editedContent: string,
+		restoreCheckpoint?: boolean,
+	): Promise<void> => {
 		// Only proceed if we have a current cline
 		if (provider.getCurrentCline()) {
 			const currentCline = provider.getCurrentCline()!
@@ -175,6 +278,22 @@ export const webviewMessageHandler = async (
 
 			if (messageIndex !== -1) {
 				try {
+					// If checkpoint restoration is requested, restore to the checkpoint first
+					if (restoreCheckpoint) {
+						const targetMessage = currentCline.clineMessages[messageIndex]
+						if (
+							targetMessage?.checkpoint &&
+							typeof targetMessage.checkpoint === "object" &&
+							"hash" in targetMessage.checkpoint
+						) {
+							await currentCline.checkpointRestore({
+								ts: targetMessage.ts!,
+								commitHash: targetMessage.checkpoint.hash as string,
+								mode: "restore",
+							})
+						}
+					}
+
 					// Edit this message and delete subsequent
 					await removeMessagesThisAndSubsequent(currentCline, messageIndex, apiConversationHistoryIndex)
 
@@ -1503,12 +1622,12 @@ export const webviewMessageHandler = async (
 			break
 		case "deleteMessageConfirm":
 			if (message.messageTs) {
-				await handleDeleteMessageConfirm(message.messageTs)
+				await handleDeleteMessageConfirm(message.messageTs, message.restoreCheckpoint)
 			}
 			break
 		case "editMessageConfirm":
 			if (message.messageTs && message.text) {
-				await handleEditMessageConfirm(message.messageTs, message.text)
+				await handleEditMessageConfirm(message.messageTs, message.text, message.restoreCheckpoint)
 			}
 			break
 		case "getListApiConfiguration":

src/shared/ExtensionMessage.ts

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

src/shared/WebviewMessage.ts

@@ -230,6 +230,7 @@ export interface WebviewMessage {
 	hasSystemPromptOverride?: boolean
 	terminalOperation?: "continue" | "abort"
 	messageTs?: number
+	restoreCheckpoint?: boolean
 	historyPreviewCollapsed?: boolean
 	filters?: { type?: string; search?: string; tags?: string[] }
 	url?: string // For openExternal

webview-ui/src/App.tsx

@@ -18,7 +18,8 @@ import McpView from "./components/mcp/McpView"
 import { MarketplaceView } from "./components/marketplace/MarketplaceView"
 import ModesView from "./components/modes/ModesView"
 import { HumanRelayDialog } from "./components/human-relay/HumanRelayDialog"
-import { DeleteMessageDialog, EditMessageDialog } from "./components/chat/MessageModificationConfirmationDialog"
+import { CheckpointRestoreDialog } from "./components/chat/CheckpointRestoreDialog"
+import { MessageModificationConfirmationDialog } from "./components/chat/MessageModificationConfirmationDialog"
 import { AccountView } from "./components/account/AccountView"
 import { useAddNonInteractiveClickListener } from "./components/ui/hooks/useNonInteractiveClick"
 import { TooltipProvider } from "./components/ui/tooltip"
@@ -74,19 +75,23 @@ const App = () => {
 	const [deleteMessageDialogState, setDeleteMessageDialogState] = useState<{
 		isOpen: boolean
 		messageTs: number
+		hasCheckpoint: boolean
 	}>({
 		isOpen: false,
 		messageTs: 0,
+		hasCheckpoint: false,
 	})
 
 	const [editMessageDialogState, setEditMessageDialogState] = useState<{
 		isOpen: boolean
 		messageTs: number
 		text: string
+		hasCheckpoint: boolean
 	}>({
 		isOpen: false,
 		messageTs: 0,
 		text: "",
+		hasCheckpoint: false,
 	})
 
 	const settingsRef = useRef<SettingsViewRef>(null)
@@ -153,7 +158,11 @@ const App = () => {
 						messageTs: message.messageTs,
 					})
 				} else {
-					setDeleteMessageDialogState({ isOpen: true, messageTs: message.messageTs })
+					setDeleteMessageDialogState({
+						isOpen: true,
+						messageTs: message.messageTs,
+						hasCheckpoint: message.hasCheckpoint || false,
+					})
 				}
 			}
 
@@ -167,7 +176,12 @@ const App = () => {
 						text: message.text,
 					})
 				} else {
-					setEditMessageDialogState({ isOpen: true, messageTs: message.messageTs, text: message.text })
+					setEditMessageDialogState({
+						isOpen: true,
+						messageTs: message.messageTs,
+						text: message.text,
+						hasCheckpoint: message.hasCheckpoint || false,
+					})
 				}
 			}
 
@@ -257,45 +271,84 @@ const App = () => {
 				onSubmit={(requestId, text) => vscode.postMessage({ type: "humanRelayResponse", requestId, text })}
 				onCancel={(requestId) => vscode.postMessage({ type: "humanRelayCancel", requestId })}
 			/>
-			<DeleteMessageDialog
+			{deleteMessageDialogState.hasCheckpoint ? (
-				open={deleteMessageDialogState.isOpen}
+				<CheckpointRestoreDialog
-				onOpenChange={(open) => setDeleteMessageDialogState((prev) => ({ ...prev, isOpen: open }))}
+					open={deleteMessageDialogState.isOpen}
-				onConfirm={(dontShowAgain) => {
+					type="delete"
-					// Save the preference if checkbox was checked
+					hasCheckpoint={deleteMessageDialogState.hasCheckpoint}
-					if (dontShowAgain) {
+					onOpenChange={(open: boolean) => setDeleteMessageDialogState((prev) => ({ ...prev, isOpen: open }))}
-						setSkipDeleteMessageConfirmation(true)
+					onConfirm={(restoreCheckpoint: boolean) => {
 						vscode.postMessage({
-							type: "skipDeleteMessageConfirmation",
+							type: "deleteMessageConfirm",
-							bool: true,
+							messageTs: deleteMessageDialogState.messageTs,
+							restoreCheckpoint,
 						})
-					}
+						setDeleteMessageDialogState((prev) => ({ ...prev, isOpen: false }))
-					vscode.postMessage({
+					}}
-						type: "deleteMessageConfirm",
+				/>
-						messageTs: deleteMessageDialogState.messageTs,
+			) : (
-					})
+				<MessageModificationConfirmationDialog
-					setDeleteMessageDialogState((prev) => ({ ...prev, isOpen: false }))
+					open={deleteMessageDialogState.isOpen}
-				}}
+					type="delete"
-			/>
+					onOpenChange={(open: boolean) => setDeleteMessageDialogState((prev) => ({ ...prev, isOpen: open }))}
-			<EditMessageDialog
+					onConfirm={(dontShowAgain: boolean) => {
-				open={editMessageDialogState.isOpen}
+						// Save the preference if checkbox was checked
-				onOpenChange={(open) => setEditMessageDialogState((prev) => ({ ...prev, isOpen: open }))}
+						if (dontShowAgain) {
-				onConfirm={(dontShowAgain) => {
+							setSkipDeleteMessageConfirmation(true)
-					// Save the preference if checkbox was checked
+							vscode.postMessage({
-					if (dontShowAgain) {
+								type: "skipDeleteMessageConfirmation",
-						setSkipEditMessageConfirmation(true)
+								bool: true,
+							})
+						}
 						vscode.postMessage({
-							type: "skipEditMessageConfirmation",
+							type: "deleteMessageConfirm",
-							bool: true,
+							messageTs: deleteMessageDialogState.messageTs,
+							restoreCheckpoint: false,
 						})
-					}
+						setDeleteMessageDialogState((prev) => ({ ...prev, isOpen: false }))
-					vscode.postMessage({
+					}}
-						type: "editMessageConfirm",
+				/>
-						messageTs: editMessageDialogState.messageTs,
+			)}
-						text: editMessageDialogState.text,
+			{editMessageDialogState.hasCheckpoint ? (
-					})
+				<CheckpointRestoreDialog
-					setEditMessageDialogState((prev) => ({ ...prev, isOpen: false }))
+					open={editMessageDialogState.isOpen}
-				}}
+					type="edit"
-			/>
+					hasCheckpoint={editMessageDialogState.hasCheckpoint}
+					onOpenChange={(open: boolean) => setEditMessageDialogState((prev) => ({ ...prev, isOpen: open }))}
+					onConfirm={(restoreCheckpoint: boolean) => {
+						vscode.postMessage({
+							type: "editMessageConfirm",
+							messageTs: editMessageDialogState.messageTs,
+							text: editMessageDialogState.text,
+							restoreCheckpoint,
+						})
+						setEditMessageDialogState((prev) => ({ ...prev, isOpen: false }))
+					}}
+				/>
+			) : (
+				<MessageModificationConfirmationDialog
+					open={editMessageDialogState.isOpen}
+					type="edit"
+					onOpenChange={(open: boolean) => setEditMessageDialogState((prev) => ({ ...prev, isOpen: open }))}
+					onConfirm={(dontShowAgain: boolean) => {
+						// Save the preference if checkbox was checked
+						if (dontShowAgain) {
+							setSkipEditMessageConfirmation(true)
+							vscode.postMessage({
+								type: "skipEditMessageConfirmation",
+								bool: true,
+							})
+						}
+						vscode.postMessage({
+							type: "editMessageConfirm",
+							messageTs: editMessageDialogState.messageTs,
+							text: editMessageDialogState.text,
+							restoreCheckpoint: false,
+						})
+						setEditMessageDialogState((prev) => ({ ...prev, isOpen: false }))
+					}}
+				/>
+			)}
 		</>
 	)
 }

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

@@ -1073,7 +1073,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/CheckpointRestoreDialog.tsx

@@ -0,0 +1,83 @@
+import React from "react"
+import { useAppTranslation } from "@src/i18n/TranslationContext"
+import {
+	AlertDialog,
+	AlertDialogAction,
+	AlertDialogCancel,
+	AlertDialogContent,
+	AlertDialogDescription,
+	AlertDialogFooter,
+	AlertDialogHeader,
+	AlertDialogTitle,
+} from "@src/components/ui"
+
+interface CheckpointRestoreDialogProps {
+	open: boolean
+	onOpenChange: (open: boolean) => void
+	onConfirm: (restoreCheckpoint: boolean) => void
+	type: "edit" | "delete"
+	hasCheckpoint: boolean
+}
+
+export const CheckpointRestoreDialog: React.FC<CheckpointRestoreDialogProps> = ({
+	open,
+	onOpenChange,
+	onConfirm,
+	type,
+	hasCheckpoint,
+}) => {
+	const { t } = useAppTranslation()
+
+	const isEdit = type === "edit"
+	const title = isEdit ? t("common:confirmation.edit_message") : t("common:confirmation.delete_message")
+	const description = isEdit
+		? t("common:confirmation.edit_question_with_checkpoint")
+		: t("common:confirmation.delete_question_with_checkpoint")
+
+	const handleConfirmWithRestore = () => {
+		onConfirm(true)
+		onOpenChange(false)
+	}
+
+	const handleConfirmWithoutRestore = () => {
+		onConfirm(false)
+		onOpenChange(false)
+	}
+
+	return (
+		<AlertDialog open={open} onOpenChange={onOpenChange}>
+			<AlertDialogContent>
+				<AlertDialogHeader>
+					<AlertDialogTitle className="text-lg">{title}</AlertDialogTitle>
+					<AlertDialogDescription className="text-base">{description}</AlertDialogDescription>
+				</AlertDialogHeader>
+				<AlertDialogFooter className="flex-col gap-2">
+					<AlertDialogCancel className="bg-vscode-button-secondaryBackground hover:bg-vscode-button-secondaryHoverBackground text-vscode-button-secondaryForeground border-vscode-button-border">
+						{t("common:answers.cancel")}
+					</AlertDialogCancel>
+					<AlertDialogAction
+						onClick={handleConfirmWithoutRestore}
+						className="bg-vscode-button-background hover:bg-vscode-button-hoverBackground text-vscode-button-foreground border-vscode-button-border">
+						{isEdit ? t("common:confirmation.edit_only") : t("common:confirmation.delete_only")}
+					</AlertDialogAction>
+					{hasCheckpoint && (
+						<AlertDialogAction
+							onClick={handleConfirmWithRestore}
+							className="bg-vscode-button-background hover:bg-vscode-button-hoverBackground text-vscode-button-foreground border-vscode-button-border">
+							{t("common:confirmation.restore_to_checkpoint")}
+						</AlertDialogAction>
+					)}
+				</AlertDialogFooter>
+			</AlertDialogContent>
+		</AlertDialog>
+	)
+}
+
+// Export convenience components for backward compatibility
+export const EditMessageWithCheckpointDialog: React.FC<Omit<CheckpointRestoreDialogProps, "type">> = (props) => (
+	<CheckpointRestoreDialog {...props} type="edit" />
+)
+
+export const DeleteMessageWithCheckpointDialog: React.FC<Omit<CheckpointRestoreDialogProps, "type">> = (props) => (
+	<CheckpointRestoreDialog {...props} type="delete" />
+)

webview-ui/src/components/chat/__tests__/CheckpointRestoreDialog.spec.tsx

@@ -0,0 +1,297 @@
+// npx vitest run src/components/chat/__tests__/CheckpointRestoreDialog.spec.tsx
+
+import React from "react"
+import { render, screen, fireEvent } from "@/utils/test-utils"
+import { vi } from "vitest"
+
+import { CheckpointRestoreDialog } from "../CheckpointRestoreDialog"
+
+// Mock the translation context
+vi.mock("@src/i18n/TranslationContext", () => ({
+	useAppTranslation: () => ({
+		t: (key: string) => {
+			const translations: Record<string, string> = {
+				"common:confirmation.delete_message": "Delete Message",
+				"common:confirmation.edit_message": "Edit Message",
+				"common:confirmation.delete_warning":
+					"Deleting this message will delete all subsequent messages in the conversation. Do you want to proceed?",
+				"common:confirmation.edit_warning":
+					"Editing this message will delete all subsequent messages in the conversation. Do you want to proceed?",
+				"common:confirmation.delete_warning_with_checkpoint":
+					"Deleting this message will delete all subsequent messages in the conversation. Do you want to proceed?",
+				"common:confirmation.edit_warning_with_checkpoint":
+					"Editing this message will delete all subsequent messages in the conversation. Do you want to proceed?",
+				"common:confirmation.restore_checkpoint": "Do you also wish to revert code to this checkpoint?",
+				"common:confirmation.proceed": "Proceed",
+				"common:answers.cancel": "Cancel",
+				"common:confirmation.dont_show_again": "Don't show this again",
+			}
+			return translations[key] || key
+		},
+	}),
+}))
+
+describe("CheckpointRestoreDialog", () => {
+	const defaultProps = {
+		open: true,
+		onOpenChange: vi.fn(),
+		onConfirm: vi.fn(),
+		type: "edit" as const,
+		hasCheckpoint: false,
+	}
+
+	beforeEach(() => {
+		vi.clearAllMocks()
+	})
+
+	describe("Basic Rendering", () => {
+		it("renders edit dialog without checkpoint", () => {
+			render(<CheckpointRestoreDialog {...defaultProps} />)
+
+			expect(screen.getByText("Edit Message")).toBeInTheDocument()
+			expect(
+				screen.getByText(
+					"Editing this message will delete all subsequent messages in the conversation. Do you want to proceed?",
+				),
+			).toBeInTheDocument()
+			expect(screen.getByText("Proceed")).toBeInTheDocument()
+			expect(screen.getByText("Cancel")).toBeInTheDocument()
+			expect(screen.queryByText("Do you also wish to revert code to this checkpoint?")).not.toBeInTheDocument()
+		})
+
+		it("renders delete dialog without checkpoint", () => {
+			render(<CheckpointRestoreDialog {...defaultProps} type="delete" />)
+
+			expect(screen.getByText("Delete Message")).toBeInTheDocument()
+			expect(
+				screen.getByText(
+					"Deleting this message will delete all subsequent messages in the conversation. Do you want to proceed?",
+				),
+			).toBeInTheDocument()
+			expect(screen.getByText("Proceed")).toBeInTheDocument()
+			expect(screen.getByText("Cancel")).toBeInTheDocument()
+			expect(screen.queryByText("Do you also wish to revert code to this checkpoint?")).not.toBeInTheDocument()
+		})
+
+		it("renders edit dialog with checkpoint option", () => {
+			render(<CheckpointRestoreDialog {...defaultProps} hasCheckpoint={true} />)
+
+			expect(screen.getByText("Edit Message")).toBeInTheDocument()
+			expect(
+				screen.getByText(
+					"Editing this message will delete all subsequent messages in the conversation. Do you want to proceed?",
+				),
+			).toBeInTheDocument()
+			expect(screen.getByText("Do you also wish to revert code to this checkpoint?")).toBeInTheDocument()
+			expect(screen.getAllByRole("checkbox")).toHaveLength(2) // restore and dont show again
+		})
+
+		it("renders delete dialog with checkpoint option", () => {
+			render(<CheckpointRestoreDialog {...defaultProps} type="delete" hasCheckpoint={true} />)
+
+			expect(screen.getByText("Delete Message")).toBeInTheDocument()
+			expect(
+				screen.getByText(
+					"Deleting this message will delete all subsequent messages in the conversation. Do you want to proceed?",
+				),
+			).toBeInTheDocument()
+			expect(screen.getByText("Do you also wish to revert code to this checkpoint?")).toBeInTheDocument()
+			expect(screen.getAllByRole("checkbox")).toHaveLength(2) // restore and dont show again
+		})
+	})
+
+	describe("User Interactions", () => {
+		it("calls onOpenChange when cancel is clicked", () => {
+			const onOpenChange = vi.fn()
+			render(<CheckpointRestoreDialog {...defaultProps} onOpenChange={onOpenChange} />)
+
+			fireEvent.click(screen.getByText("Cancel"))
+			expect(onOpenChange).toHaveBeenCalledWith(false)
+		})
+
+		it("calls onConfirm with correct parameters when proceed is clicked without checkpoint", () => {
+			const onConfirm = vi.fn()
+			render(<CheckpointRestoreDialog {...defaultProps} onConfirm={onConfirm} />)
+
+			fireEvent.click(screen.getByText("Proceed"))
+			expect(onConfirm).toHaveBeenCalledWith(false, false) // dontShowAgain, restoreCheckpoint
+		})
+
+		it("calls onConfirm with restoreCheckpoint=false when proceed is clicked with unchecked checkbox", () => {
+			const onConfirm = vi.fn()
+			render(<CheckpointRestoreDialog {...defaultProps} onConfirm={onConfirm} hasCheckpoint={true} />)
+
+			const restoreCheckbox = screen.getByLabelText("Do you also wish to revert code to this checkpoint?")
+			expect(restoreCheckbox).not.toBeChecked()
+
+			fireEvent.click(screen.getByText("Proceed"))
+			expect(onConfirm).toHaveBeenCalledWith(false, false) // dontShowAgain, restoreCheckpoint
+		})
+
+		it("calls onConfirm with restoreCheckpoint=true when proceed is clicked with checked checkbox", () => {
+			const onConfirm = vi.fn()
+			render(<CheckpointRestoreDialog {...defaultProps} onConfirm={onConfirm} hasCheckpoint={true} />)
+
+			const restoreCheckbox = screen.getByLabelText("Do you also wish to revert code to this checkpoint?")
+			fireEvent.click(restoreCheckbox)
+			expect(restoreCheckbox).toBeChecked()
+
+			fireEvent.click(screen.getByText("Proceed"))
+			expect(onConfirm).toHaveBeenCalledWith(false, true) // dontShowAgain, restoreCheckpoint
+		})
+
+		it("toggles restore checkpoint checkbox state when clicked", () => {
+			render(<CheckpointRestoreDialog {...defaultProps} hasCheckpoint={true} />)
+
+			const restoreCheckbox = screen.getByLabelText("Do you also wish to revert code to this checkpoint?")
+			expect(restoreCheckbox).not.toBeChecked()
+
+			fireEvent.click(restoreCheckbox)
+			expect(restoreCheckbox).toBeChecked()
+
+			fireEvent.click(restoreCheckbox)
+			expect(restoreCheckbox).not.toBeChecked()
+		})
+
+		it("toggles dont show again checkbox state when clicked", () => {
+			render(<CheckpointRestoreDialog {...defaultProps} hasCheckpoint={true} />)
+
+			const dontShowCheckbox = screen.getByLabelText("Don't show this again")
+			expect(dontShowCheckbox).not.toBeChecked()
+
+			fireEvent.click(dontShowCheckbox)
+			expect(dontShowCheckbox).toBeChecked()
+
+			fireEvent.click(dontShowCheckbox)
+			expect(dontShowCheckbox).not.toBeChecked()
+		})
+
+		it("calls onConfirm with dontShowAgain=true when dont show again is checked", () => {
+			const onConfirm = vi.fn()
+			render(<CheckpointRestoreDialog {...defaultProps} onConfirm={onConfirm} hasCheckpoint={true} />)
+
+			const dontShowCheckbox = screen.getByLabelText("Don't show this again")
+			fireEvent.click(dontShowCheckbox)
+			expect(dontShowCheckbox).toBeChecked()
+
+			fireEvent.click(screen.getByText("Proceed"))
+			expect(onConfirm).toHaveBeenCalledWith(true, false) // dontShowAgain, restoreCheckpoint
+		})
+	})
+
+	describe("Dialog State Management", () => {
+		it("does not render when open is false", () => {
+			render(<CheckpointRestoreDialog {...defaultProps} open={false} />)
+
+			expect(screen.queryByText("Edit Message")).not.toBeInTheDocument()
+			expect(screen.queryByText("Delete Message")).not.toBeInTheDocument()
+		})
+
+		it("resets checkbox states when dialog reopens", async () => {
+			const { rerender } = render(<CheckpointRestoreDialog {...defaultProps} hasCheckpoint={true} open={true} />)
+
+			// Check both checkboxes
+			const restoreCheckbox = screen.getByLabelText("Do you also wish to revert code to this checkpoint?")
+			const dontShowCheckbox = screen.getByLabelText("Don't show this again")
+
+			fireEvent.click(restoreCheckbox)
+			fireEvent.click(dontShowCheckbox)
+			expect(restoreCheckbox).toBeChecked()
+			expect(dontShowCheckbox).toBeChecked()
+
+			// Close dialog
+			rerender(<CheckpointRestoreDialog {...defaultProps} hasCheckpoint={true} open={false} />)
+
+			// Reopen dialog - useEffect should reset state when open becomes true
+			rerender(<CheckpointRestoreDialog {...defaultProps} hasCheckpoint={true} open={true} />)
+
+			// Checkboxes should be unchecked after reopening due to useEffect
+			const newRestoreCheckbox = screen.getByLabelText("Do you also wish to revert code to this checkpoint?")
+			const newDontShowCheckbox = screen.getByLabelText("Don't show this again")
+			expect(newRestoreCheckbox).not.toBeChecked()
+			expect(newDontShowCheckbox).not.toBeChecked()
+		})
+	})
+
+	describe("Accessibility", () => {
+		it("has proper ARIA labels and roles", () => {
+			render(<CheckpointRestoreDialog {...defaultProps} hasCheckpoint={true} />)
+
+			expect(screen.getByRole("alertdialog")).toBeInTheDocument() // AlertDialog uses alertdialog role
+			expect(screen.getAllByRole("checkbox")).toHaveLength(2) // restore and dont show again
+			expect(screen.getByRole("button", { name: "Proceed" })).toBeInTheDocument()
+			expect(screen.getByRole("button", { name: "Cancel" })).toBeInTheDocument()
+		})
+
+		it("checkboxes are properly labeled", () => {
+			render(<CheckpointRestoreDialog {...defaultProps} hasCheckpoint={true} />)
+
+			const restoreCheckbox = screen.getByLabelText("Do you also wish to revert code to this checkpoint?")
+			const dontShowCheckbox = screen.getByLabelText("Don't show this again")
+
+			expect(restoreCheckbox).toBeInTheDocument()
+			expect(dontShowCheckbox).toBeInTheDocument()
+		})
+	})
+
+	describe("Edge Cases", () => {
+		it("handles missing translation keys gracefully", () => {
+			// This test is simplified since we can't easily mock the translation function mid-test
+			// The component should handle missing keys by returning the key itself
+			render(<CheckpointRestoreDialog {...defaultProps} />)
+
+			// Should still render with proper text from our mock
+			expect(screen.getByText("Edit Message")).toBeInTheDocument()
+		})
+
+		it("handles rapid state changes", async () => {
+			const onConfirm = vi.fn()
+			render(<CheckpointRestoreDialog {...defaultProps} onConfirm={onConfirm} hasCheckpoint={true} />)
+
+			const restoreCheckbox = screen.getByLabelText("Do you also wish to revert code to this checkpoint?")
+			const proceedButton = screen.getByText("Proceed")
+
+			// Rapidly toggle checkbox and click proceed
+			fireEvent.click(restoreCheckbox)
+			fireEvent.click(restoreCheckbox)
+			fireEvent.click(restoreCheckbox)
+			fireEvent.click(proceedButton)
+
+			expect(onConfirm).toHaveBeenCalledWith(false, true) // dontShowAgain, restoreCheckpoint
+		})
+	})
+
+	describe("Type-specific Behavior", () => {
+		it("shows correct warning text for edit type", () => {
+			render(<CheckpointRestoreDialog {...defaultProps} type="edit" />)
+
+			expect(
+				screen.getByText(
+					"Editing this message will delete all subsequent messages in the conversation. Do you want to proceed?",
+				),
+			).toBeInTheDocument()
+		})
+
+		it("shows correct warning text for delete type", () => {
+			render(<CheckpointRestoreDialog {...defaultProps} type="delete" />)
+
+			expect(
+				screen.getByText(
+					"Deleting this message will delete all subsequent messages in the conversation. Do you want to proceed?",
+				),
+			).toBeInTheDocument()
+		})
+
+		it("shows correct title for edit type", () => {
+			render(<CheckpointRestoreDialog {...defaultProps} type="edit" />)
+
+			expect(screen.getByText("Edit Message")).toBeInTheDocument()
+		})
+
+		it("shows correct title for delete type", () => {
+			render(<CheckpointRestoreDialog {...defaultProps} type="delete" />)
+
+			expect(screen.getByText("Delete Message")).toBeInTheDocument()
+		})
+	})
+})

webview-ui/src/i18n/locales/en/common.json

@@ -57,6 +57,11 @@
 		"delete_warning": "Deleting this message will delete all subsequent messages in the conversation. Do you want to proceed?",
 		"edit_message": "Edit Message",
 		"edit_warning": "Editing this message will delete all subsequent messages in the conversation. Do you want to proceed?",
+		"edit_question_with_checkpoint": "Editing this message will delete all later messages in the conversation. Do you also want to undo all code changes back to this checkpoint?",
+		"delete_question_with_checkpoint": "Deleting this message will delete all later messages in the conversation. Do you also want to undo all code changes back to this checkpoint?",
+		"edit_only": "No, edit message only",
+		"delete_only": "No, delete message only",
+		"restore_to_checkpoint": "Yes, restore code to checkpoint",
 		"proceed": "Proceed",
 		"dont_show_again": "Don't show this again"
 	}