Edit/Delete User Message (#7447)

src/core/checkpoints/__tests__/checkpoint.test.ts

@@ -0,0 +1,432 @@
+import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"
+import { Task } from "../../task/Task"
+import { ClineProvider } from "../../webview/ClineProvider"
+import { checkpointSave, checkpointRestore, checkpointDiff, getCheckpointService } from "../index"
+import * as vscode from "vscode"
+
+// Mock vscode
+vi.mock("vscode", () => ({
+	window: {
+		showErrorMessage: vi.fn(),
+		createTextEditorDecorationType: vi.fn(() => ({})),
+		showInformationMessage: vi.fn(),
+	},
+	Uri: {
+		file: vi.fn((path: string) => ({ fsPath: path })),
+		parse: vi.fn((uri: string) => ({ with: vi.fn(() => ({})) })),
+	},
+	commands: {
+		executeCommand: vi.fn(),
+	},
+}))
+
+// Mock other dependencies
+vi.mock("@roo-code/telemetry", () => ({
+	TelemetryService: {
+		instance: {
+			captureCheckpointCreated: vi.fn(),
+			captureCheckpointRestored: vi.fn(),
+			captureCheckpointDiffed: vi.fn(),
+		},
+	},
+}))
+
+vi.mock("../../../utils/path", () => ({
+	getWorkspacePath: vi.fn(() => "/test/workspace"),
+}))
+
+vi.mock("../../../services/checkpoints")
+
+describe("Checkpoint functionality", () => {
+	let mockProvider: any
+	let mockTask: any
+	let mockCheckpointService: any
+
+	beforeEach(async () => {
+		// Create mock checkpoint service
+		mockCheckpointService = {
+			isInitialized: true,
+			saveCheckpoint: vi.fn().mockResolvedValue({ commit: "test-commit-hash" }),
+			restoreCheckpoint: vi.fn().mockResolvedValue(undefined),
+			getDiff: vi.fn().mockResolvedValue([]),
+			on: vi.fn(),
+			initShadowGit: vi.fn().mockResolvedValue(undefined),
+		}
+
+		// Create mock provider
+		mockProvider = {
+			context: {
+				globalStorageUri: { fsPath: "/test/storage" },
+			},
+			log: vi.fn(),
+			postMessageToWebview: vi.fn(),
+			postStateToWebview: vi.fn(),
+			cancelTask: vi.fn(),
+		}
+
+		// Create mock task
+		mockTask = {
+			taskId: "test-task-id",
+			enableCheckpoints: true,
+			checkpointService: mockCheckpointService,
+			checkpointServiceInitializing: false,
+			providerRef: {
+				deref: () => mockProvider,
+			},
+			clineMessages: [],
+			apiConversationHistory: [],
+			pendingUserMessageCheckpoint: undefined,
+			say: vi.fn().mockResolvedValue(undefined),
+			overwriteClineMessages: vi.fn(),
+			overwriteApiConversationHistory: vi.fn(),
+			combineMessages: vi.fn().mockReturnValue([]),
+		}
+
+		// Update the mock to return our mockCheckpointService
+		const checkpointsModule = await import("../../../services/checkpoints")
+		vi.mocked(checkpointsModule.RepoPerTaskCheckpointService.create).mockReturnValue(mockCheckpointService)
+	})
+
+	afterEach(() => {
+		vi.clearAllMocks()
+	})
+
+	describe("checkpointSave", () => {
+		it("should wait for checkpoint service initialization before saving", async () => {
+			// Set up task with uninitialized service
+			mockCheckpointService.isInitialized = false
+			mockTask.checkpointService = mockCheckpointService
+
+			// Simulate service initialization after a delay
+			setTimeout(() => {
+				mockCheckpointService.isInitialized = true
+			}, 100)
+
+			// Call checkpointSave
+			const savePromise = checkpointSave(mockTask, true)
+
+			// Wait for the save to complete
+			const result = await savePromise
+
+			// saveCheckpoint should have been called
+			expect(mockCheckpointService.saveCheckpoint).toHaveBeenCalledWith(
+				expect.stringContaining("Task: test-task-id"),
+				{ allowEmpty: true },
+			)
+
+			// Result should contain the commit hash
+			expect(result).toEqual({ commit: "test-commit-hash" })
+
+			// Task should still have checkpoints enabled
+			expect(mockTask.enableCheckpoints).toBe(true)
+		})
+
+		it("should handle timeout when service doesn't initialize", async () => {
+			// Service never initializes
+			mockCheckpointService.isInitialized = false
+
+			// Call checkpointSave with a task that has no checkpoint service
+			const taskWithNoService = {
+				...mockTask,
+				checkpointService: undefined,
+				enableCheckpoints: false,
+			}
+
+			const result = await checkpointSave(taskWithNoService, true)
+
+			// Result should be undefined
+			expect(result).toBeUndefined()
+
+			// saveCheckpoint should not have been called
+			expect(mockCheckpointService.saveCheckpoint).not.toHaveBeenCalled()
+		})
+
+		it("should preserve checkpoint data through message deletion flow", async () => {
+			// Initialize service
+			mockCheckpointService.isInitialized = true
+			mockTask.checkpointService = mockCheckpointService
+
+			// Simulate saving checkpoint before user message
+			const checkpointResult = await checkpointSave(mockTask, true)
+			expect(checkpointResult).toEqual({ commit: "test-commit-hash" })
+
+			// Simulate setting pendingUserMessageCheckpoint
+			if (checkpointResult && "commit" in checkpointResult) {
+				mockTask.pendingUserMessageCheckpoint = {
+					hash: checkpointResult.commit,
+					timestamp: Date.now(),
+					type: "user_message",
+				}
+			}
+
+			// Verify checkpoint data is preserved
+			expect(mockTask.pendingUserMessageCheckpoint).toBeDefined()
+			expect(mockTask.pendingUserMessageCheckpoint.hash).toBe("test-commit-hash")
+
+			// Simulate message deletion and reinitialization
+			mockTask.clineMessages = []
+			mockTask.checkpointService = mockCheckpointService // Keep service available
+			mockTask.checkpointServiceInitializing = false
+
+			// Save checkpoint again after deletion
+			const newCheckpointResult = await checkpointSave(mockTask, true)
+
+			// Should still work after reinitialization
+			expect(newCheckpointResult).toEqual({ commit: "test-commit-hash" })
+			expect(mockTask.enableCheckpoints).toBe(true)
+		})
+
+		it("should handle errors gracefully and disable checkpoints", async () => {
+			mockCheckpointService.saveCheckpoint.mockRejectedValue(new Error("Save failed"))
+
+			const result = await checkpointSave(mockTask)
+
+			expect(result).toBeUndefined()
+			expect(mockTask.enableCheckpoints).toBe(false)
+		})
+	})
+
+	describe("checkpointRestore", () => {
+		beforeEach(() => {
+			mockTask.clineMessages = [
+				{ ts: 1, say: "user", text: "Message 1" },
+				{ ts: 2, say: "assistant", text: "Message 2" },
+				{ ts: 3, say: "user", text: "Message 3" },
+			]
+			mockTask.apiConversationHistory = [
+				{ ts: 1, role: "user", content: [{ type: "text", text: "Message 1" }] },
+				{ ts: 2, role: "assistant", content: [{ type: "text", text: "Message 2" }] },
+				{ ts: 3, role: "user", content: [{ type: "text", text: "Message 3" }] },
+			]
+		})
+
+		it("should restore checkpoint for delete operation", async () => {
+			await checkpointRestore(mockTask, {
+				ts: 2,
+				commitHash: "abc123",
+				mode: "restore",
+				operation: "delete",
+			})
+
+			expect(mockCheckpointService.restoreCheckpoint).toHaveBeenCalledWith("abc123")
+			expect(mockTask.overwriteApiConversationHistory).toHaveBeenCalledWith([
+				{ ts: 1, role: "user", content: [{ type: "text", text: "Message 1" }] },
+			])
+			expect(mockTask.overwriteClineMessages).toHaveBeenCalledWith([{ ts: 1, say: "user", text: "Message 1" }])
+			expect(mockProvider.cancelTask).toHaveBeenCalled()
+		})
+
+		it("should restore checkpoint for edit operation", async () => {
+			await checkpointRestore(mockTask, {
+				ts: 2,
+				commitHash: "abc123",
+				mode: "restore",
+				operation: "edit",
+			})
+
+			expect(mockCheckpointService.restoreCheckpoint).toHaveBeenCalledWith("abc123")
+			expect(mockTask.overwriteApiConversationHistory).toHaveBeenCalledWith([
+				{ ts: 1, role: "user", content: [{ type: "text", text: "Message 1" }] },
+			])
+			// For edit operation, should include the message being edited
+			expect(mockTask.overwriteClineMessages).toHaveBeenCalledWith([
+				{ ts: 1, say: "user", text: "Message 1" },
+				{ ts: 2, say: "assistant", text: "Message 2" },
+			])
+			expect(mockProvider.cancelTask).toHaveBeenCalled()
+		})
+
+		it("should handle preview mode without modifying messages", async () => {
+			await checkpointRestore(mockTask, {
+				ts: 2,
+				commitHash: "abc123",
+				mode: "preview",
+			})
+
+			expect(mockCheckpointService.restoreCheckpoint).toHaveBeenCalledWith("abc123")
+			expect(mockTask.overwriteApiConversationHistory).not.toHaveBeenCalled()
+			expect(mockTask.overwriteClineMessages).not.toHaveBeenCalled()
+			expect(mockProvider.cancelTask).toHaveBeenCalled()
+		})
+
+		it("should handle missing message gracefully", async () => {
+			await checkpointRestore(mockTask, {
+				ts: 999, // Non-existent timestamp
+				commitHash: "abc123",
+				mode: "restore",
+			})
+
+			expect(mockCheckpointService.restoreCheckpoint).not.toHaveBeenCalled()
+		})
+
+		it("should disable checkpoints on error", async () => {
+			mockCheckpointService.restoreCheckpoint.mockRejectedValue(new Error("Restore failed"))
+
+			await checkpointRestore(mockTask, {
+				ts: 2,
+				commitHash: "abc123",
+				mode: "restore",
+			})
+
+			expect(mockTask.enableCheckpoints).toBe(false)
+			expect(mockProvider.log).toHaveBeenCalledWith("[checkpointRestore] disabling checkpoints for this task")
+		})
+	})
+
+	describe("checkpointDiff", () => {
+		beforeEach(() => {
+			mockTask.clineMessages = [
+				{ ts: 1, say: "user", text: "Message 1" },
+				{ ts: 2, say: "checkpoint_saved", text: "commit1" },
+				{ ts: 3, say: "user", text: "Message 2" },
+				{ ts: 4, say: "checkpoint_saved", text: "commit2" },
+			]
+		})
+
+		it("should show diff for full mode", async () => {
+			const mockChanges = [
+				{
+					paths: { absolute: "/test/file.ts", relative: "file.ts" },
+					content: { before: "old content", after: "new content" },
+				},
+			]
+			mockCheckpointService.getDiff.mockResolvedValue(mockChanges)
+
+			await checkpointDiff(mockTask, {
+				ts: 4,
+				commitHash: "commit2",
+				mode: "full",
+			})
+
+			expect(mockCheckpointService.getDiff).toHaveBeenCalledWith({
+				from: "commit2",
+				to: undefined,
+			})
+			expect(vscode.commands.executeCommand).toHaveBeenCalledWith(
+				"vscode.changes",
+				"Changes since task started",
+				expect.any(Array),
+			)
+		})
+
+		it("should show diff for checkpoint mode with next commit", async () => {
+			const mockChanges = [
+				{
+					paths: { absolute: "/test/file.ts", relative: "file.ts" },
+					content: { before: "old content", after: "new content" },
+				},
+			]
+			mockCheckpointService.getDiff.mockResolvedValue(mockChanges)
+			await checkpointDiff(mockTask, {
+				ts: 4,
+				commitHash: "commit1",
+				mode: "checkpoint",
+			})
+
+			expect(mockCheckpointService.getDiff).toHaveBeenCalledWith({
+				from: "commit1",
+				to: "commit2",
+			})
+			expect(vscode.commands.executeCommand).toHaveBeenCalledWith(
+				"vscode.changes",
+				"Changes compare with next checkpoint",
+				expect.any(Array),
+			)
+		})
+
+		it("should find next checkpoint automatically in checkpoint mode", async () => {
+			const mockChanges = [
+				{
+					paths: { absolute: "/test/file.ts", relative: "file.ts" },
+					content: { before: "old content", after: "new content" },
+				},
+			]
+			mockCheckpointService.getDiff.mockResolvedValue(mockChanges)
+
+			await checkpointDiff(mockTask, {
+				ts: 4,
+				commitHash: "commit1",
+				mode: "checkpoint",
+			})
+
+			expect(mockCheckpointService.getDiff).toHaveBeenCalledWith({
+				from: "commit1", // Should find the next checkpoint
+				to: "commit2",
+			})
+		})
+
+		it("should show information message when no changes found", async () => {
+			mockCheckpointService.getDiff.mockResolvedValue([])
+
+			await checkpointDiff(mockTask, {
+				ts: 4,
+				commitHash: "commit2",
+				mode: "full",
+			})
+
+			expect(vscode.window.showInformationMessage).toHaveBeenCalledWith("No changes found.")
+			expect(vscode.commands.executeCommand).not.toHaveBeenCalled()
+		})
+
+		it("should disable checkpoints on error", async () => {
+			mockCheckpointService.getDiff.mockRejectedValue(new Error("Diff failed"))
+
+			await checkpointDiff(mockTask, {
+				ts: 4,
+				commitHash: "commit2",
+				mode: "full",
+			})
+
+			expect(mockTask.enableCheckpoints).toBe(false)
+			expect(mockProvider.log).toHaveBeenCalledWith("[checkpointDiff] disabling checkpoints for this task")
+		})
+	})
+
+	describe("getCheckpointService", () => {
+		it("should return existing service if available", async () => {
+			const service = await getCheckpointService(mockTask)
+			expect(service).toBe(mockCheckpointService)
+		})
+
+		it("should return undefined if checkpoints are disabled", async () => {
+			mockTask.enableCheckpoints = false
+			const service = await getCheckpointService(mockTask)
+			expect(service).toBeUndefined()
+		})
+
+		it("should return undefined if service is still initializing", async () => {
+			mockTask.checkpointService = undefined
+			mockTask.checkpointServiceInitializing = true
+			const service = await getCheckpointService(mockTask)
+			expect(service).toBeUndefined()
+		})
+
+		it("should create new service if none exists", async () => {
+			mockTask.checkpointService = undefined
+			mockTask.checkpointServiceInitializing = false
+
+			const service = getCheckpointService(mockTask)
+
+			const checkpointsModule = await import("../../../services/checkpoints")
+			expect(vi.mocked(checkpointsModule.RepoPerTaskCheckpointService.create)).toHaveBeenCalledWith({
+				taskId: "test-task-id",
+				workspaceDir: "/test/workspace",
+				shadowDir: "/test/storage",
+				log: expect.any(Function),
+			})
+		})
+
+		it("should disable checkpoints if workspace path is not found", async () => {
+			const pathModule = await import("../../../utils/path")
+			vi.mocked(pathModule.getWorkspacePath).mockReturnValue(null as any)
+
+			mockTask.checkpointService = undefined
+			mockTask.checkpointServiceInitializing = false
+
+			const service = await getCheckpointService(mockTask)
+
+			expect(service).toBeUndefined()
+			expect(mockTask.enableCheckpoints).toBe(false)
+		})
+	})
+})

src/core/checkpoints/index.ts

@@ -184,9 +184,13 @@ export type CheckpointRestoreOptions = {
 	ts: number
 	commitHash: string
 	mode: "preview" | "restore"
+	operation?: "delete" | "edit" // Optional to maintain backward compatibility
 }
 
-export async function checkpointRestore(task: Task, { ts, commitHash, mode }: CheckpointRestoreOptions) {
+export async function checkpointRestore(
+	task: Task,
+	{ ts, commitHash, mode, operation = "delete" }: CheckpointRestoreOptions,
+) {
 	const service = await getCheckpointService(task)
 
 	if (!service) {
@@ -215,7 +219,10 @@ export async function checkpointRestore(task: Task, { ts, commitHash, mode }: Ch
 				task.combineMessages(deletedMessages),
 			)
 
-			await task.overwriteClineMessages(task.clineMessages.slice(0, index + 1))
+			// For delete operations, exclude the checkpoint message itself
+			// For edit operations, include the checkpoint message (to be edited)
+			const endIndex = operation === "edit" ? index + 1 : index
+			await task.overwriteClineMessages(task.clineMessages.slice(0, endIndex))
 
 			// TODO: Verify that this is working as expected.
 			await task.say(
@@ -264,15 +271,16 @@ export async function checkpointDiff(task: Task, { ts, previousCommitHash, commi
 	TelemetryService.instance.captureCheckpointDiffed(task.taskId)
 
 	let prevHash = commitHash
-	let nextHash: string | undefined
-
-	const checkpoints = typeof service.getCheckpoints === "function" ? service.getCheckpoints() : []
-	const idx = checkpoints.indexOf(commitHash)
-
-	if (idx !== -1 && idx < checkpoints.length - 1) {
-		nextHash = checkpoints[idx + 1]
-	} else {
-		nextHash = undefined
+	let nextHash: string | undefined = undefined
+
+	if (mode !== "full") {
+		const checkpoints = task.clineMessages.filter(({ say }) => say === "checkpoint_saved").map(({ text }) => text!)
+		const idx = checkpoints.indexOf(commitHash)
+		if (idx !== -1 && idx < checkpoints.length - 1) {
+			nextHash = checkpoints[idx + 1]
+		} else {
+			nextHash = undefined
+		}
 	}
 
 	try {
@@ -285,7 +293,7 @@ export async function checkpointDiff(task: Task, { ts, previousCommitHash, commi
 
 		await vscode.commands.executeCommand(
 			"vscode.changes",
-			mode === "full" ? "Changes since task started" : "Changes since previous checkpoint",
+			mode === "full" ? "Changes since task started" : "Changes compare with next checkpoint",
 			changes.map((change) => [
 				vscode.Uri.file(change.paths.absolute),
 				vscode.Uri.parse(`${DIFF_VIEW_URI_SCHEME}:${change.paths.relative}`).with({

src/core/webview/ClineProvider.ts

@@ -97,6 +97,20 @@ import { getUri } from "./getUri"
  * https://github.com/KumarVariable/vscode-extension-sidebar-html/blob/master/src/customSidebarViewProvider.ts
  */
 
+export type ClineProviderEvents = {
+	clineCreated: [cline: Task]
+}
+
+interface PendingEditOperation {
+	messageTs: number
+	editedContent: string
+	images?: string[]
+	messageIndex: number
+	apiConversationHistoryIndex: number
+	timeoutId: NodeJS.Timeout
+	createdAt: number
+}
+
 export class ClineProvider
 	extends EventEmitter<TaskProviderEvents>
 	implements vscode.WebviewViewProvider, TelemetryPropertiesProvider, TaskProviderLike
@@ -121,6 +135,8 @@ export class ClineProvider
 	private taskEventListeners: WeakMap<Task, Array<() => void>> = new WeakMap()
 
 	private recentTasksCache?: string[]
+	private pendingOperations: Map<string, PendingEditOperation> = new Map()
+	private static readonly PENDING_OPERATION_TIMEOUT_MS = 30000 // 30 seconds
 
 	public isViewLaunched = false
 	public settingsImportedAt?: number
@@ -440,6 +456,71 @@ export class ClineProvider
 		// the 'parent' calling task).
 		await this.getCurrentTask()?.completeSubtask(lastMessage)
 	}
+	// Pending Edit Operations Management
+
+	/**
+	 * Sets a pending edit operation with automatic timeout cleanup
+	 */
+	public setPendingEditOperation(
+		operationId: string,
+		editData: {
+			messageTs: number
+			editedContent: string
+			images?: string[]
+			messageIndex: number
+			apiConversationHistoryIndex: number
+		},
+	): void {
+		// Clear any existing operation with the same ID
+		this.clearPendingEditOperation(operationId)
+
+		// Create timeout for automatic cleanup
+		const timeoutId = setTimeout(() => {
+			this.clearPendingEditOperation(operationId)
+			this.log(`[setPendingEditOperation] Automatically cleared stale pending operation: ${operationId}`)
+		}, ClineProvider.PENDING_OPERATION_TIMEOUT_MS)
+
+		// Store the operation
+		this.pendingOperations.set(operationId, {
+			...editData,
+			timeoutId,
+			createdAt: Date.now(),
+		})
+
+		this.log(`[setPendingEditOperation] Set pending operation: ${operationId}`)
+	}
+
+	/**
+	 * Gets a pending edit operation by ID
+	 */
+	private getPendingEditOperation(operationId: string): PendingEditOperation | undefined {
+		return this.pendingOperations.get(operationId)
+	}
+
+	/**
+	 * Clears a specific pending edit operation
+	 */
+	private clearPendingEditOperation(operationId: string): boolean {
+		const operation = this.pendingOperations.get(operationId)
+		if (operation) {
+			clearTimeout(operation.timeoutId)
+			this.pendingOperations.delete(operationId)
+			this.log(`[clearPendingEditOperation] Cleared pending operation: ${operationId}`)
+			return true
+		}
+		return false
+	}
+
+	/**
+	 * Clears all pending edit operations
+	 */
+	private clearAllPendingEditOperations(): void {
+		for (const [operationId, operation] of this.pendingOperations) {
+			clearTimeout(operation.timeoutId)
+		}
+		this.pendingOperations.clear()
+		this.log(`[clearAllPendingEditOperations] Cleared all pending operations`)
+	}
 
 	/*
 	VSCode extensions use the disposable pattern to clean up resources when the sidebar/editor tab is closed by the user or system. This applies to event listening, commands, interacting with the UI, etc.
@@ -465,6 +546,10 @@ export class ClineProvider
 
 		this.log("Cleared all tasks")
 
+		// Clear all pending edit operations to prevent memory leaks
+		this.clearAllPendingEditOperations()
+		this.log("Cleared pending operations")
+
 		if (this.view && "dispose" in this.view) {
 			this.view.dispose()
 			this.log("Disposed webview")
@@ -805,6 +890,49 @@ export class ClineProvider
 			`[createTaskWithHistoryItem] ${task.parentTask ? "child" : "parent"} task ${task.taskId}.${task.instanceId} instantiated`,
 		)
 
+		// Check if there's a pending edit after checkpoint restoration
+		const operationId = `task-${task.taskId}`
+		const pendingEdit = this.getPendingEditOperation(operationId)
+		if (pendingEdit) {
+			this.clearPendingEditOperation(operationId) // Clear the pending edit
+
+			this.log(`[createTaskWithHistoryItem] Processing pending edit after checkpoint restoration`)
+
+			// Process the pending edit after a short delay to ensure the task is fully initialized
+			setTimeout(async () => {
+				try {
+					// Find the message index in the restored state
+					const { messageIndex, apiConversationHistoryIndex } = (() => {
+						const messageIndex = task.clineMessages.findIndex((msg) => msg.ts === pendingEdit.messageTs)
+						const apiConversationHistoryIndex = task.apiConversationHistory.findIndex(
+							(msg) => msg.ts === pendingEdit.messageTs,
+						)
+						return { messageIndex, apiConversationHistoryIndex }
+					})()
+
+					if (messageIndex !== -1) {
+						// Remove the target message and all subsequent messages
+						await task.overwriteClineMessages(task.clineMessages.slice(0, messageIndex))
+
+						if (apiConversationHistoryIndex !== -1) {
+							await task.overwriteApiConversationHistory(
+								task.apiConversationHistory.slice(0, apiConversationHistoryIndex),
+							)
+						}
+
+						// Process the edited message
+						await task.handleWebviewAskResponse(
+							"messageResponse",
+							pendingEdit.editedContent,
+							pendingEdit.images,
+						)
+					}
+				} catch (error) {
+					this.log(`[createTaskWithHistoryItem] Error processing pending edit: ${error}`)
+				}
+			}, 100) // Small delay to ensure task is fully ready
+		}
+
 		return task
 	}
 

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

@@ -44,6 +44,12 @@ vi.mock("axios", () => ({
 
 vi.mock("../../../utils/safeWriteJson")
 
+vi.mock("../../../utils/storage", () => ({
+	getSettingsDirectoryPath: vi.fn().mockResolvedValue("/test/settings/path"),
+	getTaskDirectoryPath: vi.fn().mockResolvedValue("/test/task/path"),
+	getGlobalStoragePath: vi.fn().mockResolvedValue("/test/storage/path"),
+}))
+
 vi.mock("@modelcontextprotocol/sdk/types.js", () => ({
 	CallToolResultSchema: {},
 	ListResourcesResultSchema: {},
@@ -1179,8 +1185,8 @@ describe("ClineProvider", () => {
 			const mockMessages = [
 				{ ts: 1000, type: "say", say: "user_feedback" }, // User message 1
 				{ ts: 2000, type: "say", say: "tool" }, // Tool message
-				{ ts: 3000, type: "say", say: "text", value: 4000 }, // Message to delete
-				{ ts: 4000, type: "say", say: "browser_action" }, // Response to delete
+				{ ts: 3000, type: "say", say: "text" }, // Message before delete
+				{ ts: 4000, type: "say", say: "browser_action" }, // Message to delete
 				{ ts: 5000, type: "say", say: "user_feedback" }, // Next user message
 				{ ts: 6000, type: "say", say: "user_feedback" }, // Final message
 			] as ClineMessage[]
@@ -1216,22 +1222,28 @@ describe("ClineProvider", () => {
 			expect(mockPostMessage).toHaveBeenCalledWith({
 				type: "showDeleteMessageDialog",
 				messageTs: 4000,
+				hasCheckpoint: false,
 			})
 
 			// Simulate user confirming deletion through the dialog
 			await messageHandler({ type: "deleteMessageConfirm", messageTs: 4000 })
 
 			// Verify only messages before the deleted message were kept
-			expect(mockCline.overwriteClineMessages).toHaveBeenCalledWith([mockMessages[0], mockMessages[1]])
+			expect(mockCline.overwriteClineMessages).toHaveBeenCalledWith([
+				mockMessages[0],
+				mockMessages[1],
+				mockMessages[2],
+			])
 
 			// Verify only API messages before the deleted message were kept
 			expect(mockCline.overwriteApiConversationHistory).toHaveBeenCalledWith([
 				mockApiHistory[0],
 				mockApiHistory[1],
+				mockApiHistory[2],
 			])
 
-			// Verify createTaskWithHistoryItem was called
-			expect((provider as any).createTaskWithHistoryItem).toHaveBeenCalledWith({ id: "test-task-id" })
+			// createTaskWithHistoryItem is only called when restoring checkpoints or aborting tasks
+			expect((provider as any).createTaskWithHistoryItem).not.toHaveBeenCalled()
 		})
 
 		test("handles case when no current task exists", async () => {
@@ -1261,8 +1273,8 @@ describe("ClineProvider", () => {
 			const mockMessages = [
 				{ ts: 1000, type: "say", say: "user_feedback" }, // User message 1
 				{ ts: 2000, type: "say", say: "tool" }, // Tool message
-				{ ts: 3000, type: "say", say: "text", value: 4000 }, // Message to edit
-				{ ts: 4000, type: "say", say: "browser_action" }, // Response to edit
+				{ ts: 3000, type: "say", say: "text" }, // Message before edit
+				{ ts: 4000, type: "say", say: "browser_action" }, // Message to edit
 				{ ts: 5000, type: "say", say: "user_feedback" }, // Next user message
 				{ ts: 6000, type: "say", say: "user_feedback" }, // Final message
 			] as ClineMessage[]
@@ -1309,6 +1321,8 @@ describe("ClineProvider", () => {
 				type: "showEditMessageDialog",
 				messageTs: 4000,
 				text: "Edited message content",
+				hasCheckpoint: false,
+				images: undefined,
 			})
 
 			// Simulate user confirming edit through the dialog
@@ -1319,12 +1333,17 @@ describe("ClineProvider", () => {
 			})
 
 			// Verify correct messages were kept (only messages before the edited one)
-			expect(mockCline.overwriteClineMessages).toHaveBeenCalledWith([mockMessages[0], mockMessages[1]])
+			expect(mockCline.overwriteClineMessages).toHaveBeenCalledWith([
+				mockMessages[0],
+				mockMessages[1],
+				mockMessages[2],
+			])
 
 			// Verify correct API messages were kept (only messages before the edited one)
 			expect(mockCline.overwriteApiConversationHistory).toHaveBeenCalledWith([
 				mockApiHistory[0],
 				mockApiHistory[1],
+				mockApiHistory[2],
 			])
 
 			// The new flow calls webviewMessageHandler recursively with askResponse
@@ -3016,6 +3035,8 @@ describe("ClineProvider - Comprehensive Edit/Delete Edge Cases", () => {
 				type: "showEditMessageDialog",
 				messageTs: 3000,
 				text: "Edited message with preserved images",
+				hasCheckpoint: false,
+				images: undefined,
 			})
 
 			// Simulate confirmation
@@ -3025,9 +3046,9 @@ describe("ClineProvider - Comprehensive Edit/Delete Edge Cases", () => {
 				text: "Edited message with preserved images",
 			})
 
-			// Verify messages were edited correctly - only the first message should remain
-			expect(mockCline.overwriteClineMessages).toHaveBeenCalledWith([mockMessages[0]])
-			expect(mockCline.overwriteApiConversationHistory).toHaveBeenCalledWith([{ ts: 1000 }])
+			// Verify messages were edited correctly - messages up to the edited message should remain
+			expect(mockCline.overwriteClineMessages).toHaveBeenCalledWith([mockMessages[0], mockMessages[1]])
+			expect(mockCline.overwriteApiConversationHistory).toHaveBeenCalledWith([{ ts: 1000 }, { ts: 2000 }])
 		})
 
 		test("handles editing messages with file attachments", async () => {
@@ -3068,6 +3089,8 @@ describe("ClineProvider - Comprehensive Edit/Delete Edge Cases", () => {
 				type: "showEditMessageDialog",
 				messageTs: 3000,
 				text: "Edited message with file attachment",
+				hasCheckpoint: false,
+				images: undefined,
 			})
 
 			// Simulate user confirming the edit
@@ -3124,6 +3147,8 @@ describe("ClineProvider - Comprehensive Edit/Delete Edge Cases", () => {
 				type: "showEditMessageDialog",
 				messageTs: 2000,
 				text: "Edited message",
+				hasCheckpoint: false,
+				images: undefined,
 			})
 
 			// Simulate user confirming the edit
@@ -3164,6 +3189,8 @@ describe("ClineProvider - Comprehensive Edit/Delete Edge Cases", () => {
 				type: "showEditMessageDialog",
 				messageTs: 2000,
 				text: "Edited message",
+				hasCheckpoint: false,
+				images: undefined,
 			})
 
 			// Simulate user confirming the edit
@@ -3220,11 +3247,15 @@ describe("ClineProvider - Comprehensive Edit/Delete Edge Cases", () => {
 				type: "showEditMessageDialog",
 				messageTs: 2000,
 				text: "Edited message 1",
+				hasCheckpoint: false,
+				images: undefined,
 			})
 			expect(mockPostMessage).toHaveBeenCalledWith({
 				type: "showEditMessageDialog",
 				messageTs: 4000,
 				text: "Edited message 2",
+				hasCheckpoint: false,
+				images: undefined,
 			})
 
 			// Simulate user confirming both edits
@@ -3410,6 +3441,8 @@ describe("ClineProvider - Comprehensive Edit/Delete Edge Cases", () => {
 					type: "showEditMessageDialog",
 					messageTs: 5000,
 					text: "Edited non-existent message",
+					hasCheckpoint: false,
+					images: undefined,
 				})
 
 				// Simulate user confirming the edit
@@ -3450,6 +3483,7 @@ describe("ClineProvider - Comprehensive Edit/Delete Edge Cases", () => {
 				expect(mockPostMessage).toHaveBeenCalledWith({
 					type: "showDeleteMessageDialog",
 					messageTs: 5000,
+					hasCheckpoint: false,
 				})
 
 				// Simulate user confirming the delete
@@ -3501,6 +3535,8 @@ describe("ClineProvider - Comprehensive Edit/Delete Edge Cases", () => {
 					type: "showEditMessageDialog",
 					messageTs: 2000,
 					text: "Edited message",
+					hasCheckpoint: false,
+					images: undefined,
 				})
 
 				// Simulate user confirming the edit
@@ -3540,6 +3576,7 @@ describe("ClineProvider - Comprehensive Edit/Delete Edge Cases", () => {
 				expect(mockPostMessage).toHaveBeenCalledWith({
 					type: "showDeleteMessageDialog",
 					messageTs: 2000,
+					hasCheckpoint: false,
 				})
 
 				// Simulate user confirming the delete
@@ -3593,6 +3630,8 @@ describe("ClineProvider - Comprehensive Edit/Delete Edge Cases", () => {
 					type: "showEditMessageDialog",
 					messageTs: 2000,
 					text: largeEditedContent,
+					hasCheckpoint: false,
+					images: undefined,
 				})
 
 				// Simulate user confirming the edit
@@ -3635,18 +3674,23 @@ describe("ClineProvider - Comprehensive Edit/Delete Edge Cases", () => {
 				expect(mockPostMessage).toHaveBeenCalledWith({
 					type: "showDeleteMessageDialog",
 					messageTs: 3000,
+					hasCheckpoint: false,
 				})
 
 				// Simulate user confirming the delete
 				await messageHandler({ type: "deleteMessageConfirm", messageTs: 3000 })
 
-				// Should handle large payloads without issues
-				expect(mockCline.overwriteClineMessages).toHaveBeenCalledWith([mockMessages[0]])
-				expect(mockCline.overwriteApiConversationHistory).toHaveBeenCalledWith([{ ts: 1000 }])
+				// Should handle large payloads without issues - keeps messages before the deleted one
+				expect(mockCline.overwriteClineMessages).toHaveBeenCalledWith([mockMessages[0], mockMessages[1]])
+				expect(mockCline.overwriteApiConversationHistory).toHaveBeenCalledWith([{ ts: 1000 }, { ts: 2000 }])
 			})
 		})
 
 		describe("Error Messaging and User Feedback", () => {
+			beforeEach(async () => {
+				await provider.resolveWebviewView(mockWebviewView)
+			})
+
 			// Note: Error messaging test removed as the implementation may not have proper error handling in place
 
 			test("provides user feedback for successful operations", async () => {
@@ -3673,6 +3717,7 @@ describe("ClineProvider - Comprehensive Edit/Delete Edge Cases", () => {
 				expect(mockPostMessage).toHaveBeenCalledWith({
 					type: "showDeleteMessageDialog",
 					messageTs: 2000,
+					hasCheckpoint: false,
 				})
 
 				// Simulate user confirming the delete
@@ -3680,7 +3725,7 @@ describe("ClineProvider - Comprehensive Edit/Delete Edge Cases", () => {
 
 				// Verify successful operation completed
 				expect(mockCline.overwriteClineMessages).toHaveBeenCalled()
-				expect(provider.createTaskWithHistoryItem).toHaveBeenCalled()
+				// createTaskWithHistoryItem is only called when restoring checkpoints or aborting tasks
 				expect(vscode.window.showErrorMessage).not.toHaveBeenCalled()
 			})
 
@@ -3746,6 +3791,7 @@ describe("ClineProvider - Comprehensive Edit/Delete Edge Cases", () => {
 				expect(mockPostMessage).toHaveBeenCalledWith({
 					type: "showDeleteMessageDialog",
 					messageTs: 1000,
+					hasCheckpoint: false,
 				})
 
 				// Simulate user confirming the delete
@@ -3796,6 +3842,8 @@ describe("ClineProvider - Comprehensive Edit/Delete Edge Cases", () => {
 					type: "showEditMessageDialog",
 					messageTs: futureTimestamp + 1000,
 					text: "Edited future message",
+					hasCheckpoint: false,
+					images: undefined,
 				})
 
 				// Simulate user confirming the edit

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

@@ -0,0 +1,242 @@
+import { describe, it, expect, vi, beforeEach } from "vitest"
+import { handleCheckpointRestoreOperation } from "../checkpointRestoreHandler"
+import { saveTaskMessages } from "../../task-persistence"
+import pWaitFor from "p-wait-for"
+import * as vscode from "vscode"
+
+// Mock dependencies
+vi.mock("../../task-persistence", () => ({
+	saveTaskMessages: vi.fn(),
+}))
+vi.mock("p-wait-for")
+vi.mock("vscode", () => ({
+	window: {
+		showErrorMessage: vi.fn(),
+	},
+}))
+
+describe("checkpointRestoreHandler", () => {
+	let mockProvider: any
+	let mockCline: any
+
+	beforeEach(() => {
+		vi.clearAllMocks()
+
+		// Setup mock Cline instance
+		mockCline = {
+			taskId: "test-task-123",
+			abort: false,
+			abortTask: vi.fn(() => {
+				mockCline.abort = true
+			}),
+			checkpointRestore: vi.fn(),
+			clineMessages: [
+				{ ts: 1, type: "user", say: "user", text: "First message" },
+				{ ts: 2, type: "assistant", say: "assistant", text: "Response" },
+				{
+					ts: 3,
+					type: "user",
+					say: "user",
+					text: "Checkpoint message",
+					checkpoint: { hash: "abc123" },
+				},
+				{ ts: 4, type: "assistant", say: "assistant", text: "After checkpoint" },
+			],
+		}
+
+		// Setup mock provider
+		mockProvider = {
+			getCurrentTask: vi.fn(() => mockCline),
+			postMessageToWebview: vi.fn(),
+			getTaskWithId: vi.fn(() => ({
+				historyItem: { id: "test-task-123", messages: mockCline.clineMessages },
+			})),
+			createTaskWithHistoryItem: vi.fn(),
+			setPendingEditOperation: vi.fn(),
+			contextProxy: {
+				globalStorageUri: { fsPath: "/test/storage" },
+			},
+		}
+
+		// Mock pWaitFor to resolve immediately
+		;(pWaitFor as any).mockImplementation(async (condition: () => boolean) => {
+			// Simulate the condition being met
+			return Promise.resolve()
+		})
+	})
+
+	describe("handleCheckpointRestoreOperation", () => {
+		it("should abort task before checkpoint restore for delete operations", async () => {
+			// Simulate a task that hasn't been aborted yet
+			mockCline.abort = false
+
+			await handleCheckpointRestoreOperation({
+				provider: mockProvider,
+				currentCline: mockCline,
+				messageTs: 3,
+				messageIndex: 2,
+				checkpoint: { hash: "abc123" },
+				operation: "delete",
+			})
+
+			// Verify abortTask was called before checkpointRestore
+			expect(mockCline.abortTask).toHaveBeenCalled()
+			expect(mockCline.checkpointRestore).toHaveBeenCalled()
+
+			// Verify the order of operations
+			const abortOrder = mockCline.abortTask.mock.invocationCallOrder[0]
+			const restoreOrder = mockCline.checkpointRestore.mock.invocationCallOrder[0]
+			expect(abortOrder).toBeLessThan(restoreOrder)
+		})
+
+		it("should not abort task if already aborted", async () => {
+			// Simulate a task that's already aborted
+			mockCline.abort = true
+
+			await handleCheckpointRestoreOperation({
+				provider: mockProvider,
+				currentCline: mockCline,
+				messageTs: 3,
+				messageIndex: 2,
+				checkpoint: { hash: "abc123" },
+				operation: "delete",
+			})
+
+			// Verify abortTask was not called
+			expect(mockCline.abortTask).not.toHaveBeenCalled()
+			expect(mockCline.checkpointRestore).toHaveBeenCalled()
+		})
+
+		it("should handle edit operations with pending edit data", async () => {
+			const editData = {
+				editedContent: "Edited content",
+				images: ["image1.png"],
+				apiConversationHistoryIndex: 2,
+			}
+
+			await handleCheckpointRestoreOperation({
+				provider: mockProvider,
+				currentCline: mockCline,
+				messageTs: 3,
+				messageIndex: 2,
+				checkpoint: { hash: "abc123" },
+				operation: "edit",
+				editData,
+			})
+
+			// Verify abortTask was NOT called for edit operations
+			expect(mockCline.abortTask).not.toHaveBeenCalled()
+
+			// Verify pending edit operation was set
+			expect(mockProvider.setPendingEditOperation).toHaveBeenCalledWith("task-test-task-123", {
+				messageTs: 3,
+				editedContent: "Edited content",
+				images: ["image1.png"],
+				messageIndex: 2,
+				apiConversationHistoryIndex: 2,
+			})
+
+			// Verify checkpoint restore was called with edit operation
+			expect(mockCline.checkpointRestore).toHaveBeenCalledWith({
+				ts: 3,
+				commitHash: "abc123",
+				mode: "restore",
+				operation: "edit",
+			})
+		})
+
+		it("should save messages after delete operation", async () => {
+			// Mock the checkpoint restore to simulate message deletion
+			mockCline.checkpointRestore.mockImplementation(async () => {
+				mockCline.clineMessages = mockCline.clineMessages.slice(0, 2)
+			})
+
+			await handleCheckpointRestoreOperation({
+				provider: mockProvider,
+				currentCline: mockCline,
+				messageTs: 3,
+				messageIndex: 2,
+				checkpoint: { hash: "abc123" },
+				operation: "delete",
+			})
+
+			// Verify saveTaskMessages was called
+			expect(saveTaskMessages).toHaveBeenCalledWith({
+				messages: mockCline.clineMessages,
+				taskId: "test-task-123",
+				globalStoragePath: "/test/storage",
+			})
+
+			// Verify createTaskWithHistoryItem was called
+			expect(mockProvider.createTaskWithHistoryItem).toHaveBeenCalled()
+		})
+
+		it("should reinitialize task with correct history item after delete", async () => {
+			const expectedHistoryItem = {
+				id: "test-task-123",
+				messages: mockCline.clineMessages,
+			}
+
+			await handleCheckpointRestoreOperation({
+				provider: mockProvider,
+				currentCline: mockCline,
+				messageTs: 3,
+				messageIndex: 2,
+				checkpoint: { hash: "abc123" },
+				operation: "delete",
+			})
+
+			// Verify getTaskWithId was called
+			expect(mockProvider.getTaskWithId).toHaveBeenCalledWith("test-task-123")
+
+			// Verify createTaskWithHistoryItem was called with the correct history item
+			expect(mockProvider.createTaskWithHistoryItem).toHaveBeenCalledWith(expectedHistoryItem)
+		})
+
+		it("should not save messages or reinitialize for edit operation", async () => {
+			const editData = {
+				editedContent: "Edited content",
+				images: [],
+				apiConversationHistoryIndex: 2,
+			}
+
+			await handleCheckpointRestoreOperation({
+				provider: mockProvider,
+				currentCline: mockCline,
+				messageTs: 3,
+				messageIndex: 2,
+				checkpoint: { hash: "abc123" },
+				operation: "edit",
+				editData,
+			})
+
+			// Verify saveTaskMessages was NOT called for edit operation
+			expect(saveTaskMessages).not.toHaveBeenCalled()
+
+			// Verify createTaskWithHistoryItem was NOT called for edit operation
+			expect(mockProvider.createTaskWithHistoryItem).not.toHaveBeenCalled()
+		})
+
+		it("should handle errors gracefully", async () => {
+			// Mock checkpoint restore to throw an error
+			mockCline.checkpointRestore.mockRejectedValue(new Error("Checkpoint restore failed"))
+
+			// The function should throw and show an error message
+			await expect(
+				handleCheckpointRestoreOperation({
+					provider: mockProvider,
+					currentCline: mockCline,
+					messageTs: 3,
+					messageIndex: 2,
+					checkpoint: { hash: "abc123" },
+					operation: "delete",
+				}),
+			).rejects.toThrow("Checkpoint restore failed")
+
+			// Verify error message was shown
+			expect(vscode.window.showErrorMessage).toHaveBeenCalledWith(
+				"Error during checkpoint restore: Checkpoint restore failed",
+			)
+		})
+	})
+})

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

@@ -0,0 +1,131 @@
+import { describe, it, expect, vi, beforeEach } from "vitest"
+import { webviewMessageHandler } from "../webviewMessageHandler"
+import { saveTaskMessages } from "../../task-persistence"
+import { handleCheckpointRestoreOperation } from "../checkpointRestoreHandler"
+
+// Mock dependencies
+vi.mock("../../task-persistence")
+vi.mock("../checkpointRestoreHandler")
+vi.mock("vscode", () => ({
+	window: {
+		showErrorMessage: vi.fn(),
+	},
+	workspace: {
+		workspaceFolders: undefined,
+	},
+}))
+
+describe("webviewMessageHandler - checkpoint operations", () => {
+	let mockProvider: any
+	let mockCline: any
+
+	beforeEach(() => {
+		vi.clearAllMocks()
+
+		// Setup mock Cline instance
+		mockCline = {
+			taskId: "test-task-123",
+			clineMessages: [
+				{ ts: 1, type: "user", say: "user", text: "First message" },
+				{ ts: 2, type: "assistant", say: "checkpoint_saved", text: "abc123" },
+				{ ts: 3, type: "user", say: "user", text: "Message to delete" },
+				{ ts: 4, type: "assistant", say: "assistant", text: "After message" },
+			],
+			apiConversationHistory: [
+				{ ts: 1, role: "user", content: [{ type: "text", text: "First message" }] },
+				{ ts: 3, role: "user", content: [{ type: "text", text: "Message to delete" }] },
+				{ ts: 4, role: "assistant", content: [{ type: "text", text: "After message" }] },
+			],
+			checkpointRestore: vi.fn(),
+			overwriteClineMessages: vi.fn(),
+			overwriteApiConversationHistory: vi.fn(),
+		}
+
+		// Setup mock provider
+		mockProvider = {
+			getCurrentTask: vi.fn(() => mockCline),
+			postMessageToWebview: vi.fn(),
+			getTaskWithId: vi.fn(() => ({
+				historyItem: { id: "test-task-123", messages: mockCline.clineMessages },
+			})),
+			createTaskWithHistoryItem: vi.fn(),
+			setPendingEditOperation: vi.fn(),
+			contextProxy: {
+				globalStorageUri: { fsPath: "/test/storage" },
+			},
+		}
+	})
+
+	describe("delete operations with checkpoint restoration", () => {
+		it("should call handleCheckpointRestoreOperation for checkpoint deletes", async () => {
+			// Mock handleCheckpointRestoreOperation
+			;(handleCheckpointRestoreOperation as any).mockResolvedValue(undefined)
+
+			// Call the handler with delete confirmation
+			await webviewMessageHandler(mockProvider, {
+				type: "deleteMessageConfirm",
+				messageTs: 1,
+				restoreCheckpoint: true,
+			})
+
+			// Verify handleCheckpointRestoreOperation was called with correct parameters
+			expect(handleCheckpointRestoreOperation).toHaveBeenCalledWith({
+				provider: mockProvider,
+				currentCline: mockCline,
+				messageTs: 1,
+				messageIndex: 0,
+				checkpoint: { hash: "abc123" },
+				operation: "delete",
+			})
+		})
+
+		it("should save messages for non-checkpoint deletes", async () => {
+			// Call the handler with delete confirmation (no checkpoint restoration)
+			await webviewMessageHandler(mockProvider, {
+				type: "deleteMessageConfirm",
+				messageTs: 2,
+				restoreCheckpoint: false,
+			})
+
+			// Verify saveTaskMessages was called
+			expect(saveTaskMessages).toHaveBeenCalledWith({
+				messages: expect.any(Array),
+				taskId: "test-task-123",
+				globalStoragePath: "/test/storage",
+			})
+
+			// Verify checkpoint restore was NOT called
+			expect(mockCline.checkpointRestore).not.toHaveBeenCalled()
+		})
+	})
+
+	describe("edit operations with checkpoint restoration", () => {
+		it("should call handleCheckpointRestoreOperation for checkpoint edits", async () => {
+			// Mock handleCheckpointRestoreOperation
+			;(handleCheckpointRestoreOperation as any).mockResolvedValue(undefined)
+
+			// Call the handler with edit confirmation
+			await webviewMessageHandler(mockProvider, {
+				type: "editMessageConfirm",
+				messageTs: 1,
+				text: "Edited checkpoint message",
+				restoreCheckpoint: true,
+			})
+
+			// Verify handleCheckpointRestoreOperation was called with correct parameters
+			expect(handleCheckpointRestoreOperation).toHaveBeenCalledWith({
+				provider: mockProvider,
+				currentCline: mockCline,
+				messageTs: 1,
+				messageIndex: 0,
+				checkpoint: { hash: "abc123" },
+				operation: "edit",
+				editData: {
+					editedContent: "Edited checkpoint message",
+					images: undefined,
+					apiConversationHistoryIndex: 0,
+				},
+			})
+		})
+	})
+})

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

@@ -563,7 +563,10 @@ describe("webviewMessageHandler - message dialog preferences", () => {
 
 	describe("deleteMessage", () => {
 		it("should always show dialog for delete confirmation", async () => {
-			vi.mocked(mockClineProvider.getCurrentTask).mockReturnValue({} as any)
+			vi.mocked(mockClineProvider.getCurrentTask).mockReturnValue({
+				clineMessages: [],
+				apiConversationHistory: [],
+			} as any) // Mock current cline with proper structure
 
 			await webviewMessageHandler(mockClineProvider, {
 				type: "deleteMessage",
@@ -573,13 +576,17 @@ describe("webviewMessageHandler - message dialog preferences", () => {
 			expect(mockClineProvider.postMessageToWebview).toHaveBeenCalledWith({
 				type: "showDeleteMessageDialog",
 				messageTs: 123456789,
+				hasCheckpoint: false,
 			})
 		})
 	})
 
 	describe("submitEditedMessage", () => {
 		it("should always show dialog for edit confirmation", async () => {
-			vi.mocked(mockClineProvider.getCurrentTask).mockReturnValue({} as any)
+			vi.mocked(mockClineProvider.getCurrentTask).mockReturnValue({
+				clineMessages: [],
+				apiConversationHistory: [],
+			} as any) // Mock current cline with proper structure
 
 			await webviewMessageHandler(mockClineProvider, {
 				type: "submitEditedMessage",
@@ -591,6 +598,8 @@ describe("webviewMessageHandler - message dialog preferences", () => {
 				type: "showEditMessageDialog",
 				messageTs: 123456789,
 				text: "edited content",
+				hasCheckpoint: false,
+				images: undefined,
 			})
 		})
 	})

src/core/webview/checkpointRestoreHandler.ts

@@ -0,0 +1,104 @@
+import { Task } from "../task/Task"
+import { ClineProvider } from "./ClineProvider"
+import { saveTaskMessages } from "../task-persistence"
+import * as vscode from "vscode"
+import pWaitFor from "p-wait-for"
+import { t } from "../../i18n"
+
+export interface CheckpointRestoreConfig {
+	provider: ClineProvider
+	currentCline: Task
+	messageTs: number
+	messageIndex: number
+	checkpoint: { hash: string }
+	operation: "delete" | "edit"
+	editData?: {
+		editedContent: string
+		images?: string[]
+		apiConversationHistoryIndex: number
+	}
+}
+
+/**
+ * Handles checkpoint restoration for both delete and edit operations.
+ * This consolidates the common logic while handling operation-specific behavior.
+ */
+export async function handleCheckpointRestoreOperation(config: CheckpointRestoreConfig): Promise<void> {
+	const { provider, currentCline, messageTs, checkpoint, operation, editData } = config
+
+	try {
+		// For delete operations, ensure the task is properly aborted to handle any pending ask operations
+		// This prevents "Current ask promise was ignored" errors
+		// For edit operations, we don't abort because the checkpoint restore will handle it
+		if (operation === "delete" && currentCline && !currentCline.abort) {
+			currentCline.abortTask()
+			// Wait a bit for the abort to complete
+			await pWaitFor(() => currentCline.abort === true, {
+				timeout: 1000,
+				interval: 50,
+			}).catch(() => {
+				// Continue even if timeout - the abort flag should be set
+			})
+		}
+
+		// For edit operations, set up pending edit data before restoration
+		if (operation === "edit" && editData) {
+			const operationId = `task-${currentCline.taskId}`
+			provider.setPendingEditOperation(operationId, {
+				messageTs,
+				editedContent: editData.editedContent,
+				images: editData.images,
+				messageIndex: config.messageIndex,
+				apiConversationHistoryIndex: editData.apiConversationHistoryIndex,
+			})
+		}
+
+		// Perform the checkpoint restoration
+		await currentCline.checkpointRestore({
+			ts: messageTs,
+			commitHash: checkpoint.hash,
+			mode: "restore",
+			operation,
+		})
+
+		// For delete operations, we need to save messages and reinitialize
+		// For edit operations, the reinitialization happens automatically
+		// and processes the pending edit
+		if (operation === "delete") {
+			// Save the updated messages to disk after checkpoint restoration
+			await saveTaskMessages({
+				messages: currentCline.clineMessages,
+				taskId: currentCline.taskId,
+				globalStoragePath: provider.contextProxy.globalStorageUri.fsPath,
+			})
+
+			// Get the updated history item and reinitialize
+			const { historyItem } = await provider.getTaskWithId(currentCline.taskId)
+			await provider.createTaskWithHistoryItem(historyItem)
+		}
+		// For edit operations, the task cancellation in checkpointRestore
+		// will trigger reinitialization, which will process pendingEditAfterRestore
+	} catch (error) {
+		console.error(`Error in checkpoint restore (${operation}):`, error)
+		vscode.window.showErrorMessage(
+			`Error during checkpoint restore: ${error instanceof Error ? error.message : String(error)}`,
+		)
+		throw error
+	}
+}
+
+/**
+ * Common checkpoint restore validation and initialization utility.
+ * This can be used by any checkpoint restore flow that needs to wait for initialization.
+ */
+export async function waitForClineInitialization(provider: ClineProvider, timeoutMs: number = 3000): Promise<boolean> {
+	try {
+		await pWaitFor(() => provider.getCurrentTask()?.isInitialized === true, {
+			timeout: timeoutMs,
+		})
+		return true
+	} catch (error) {
+		vscode.window.showErrorMessage(t("common:errors.checkpoint_timeout"))
+		return false
+	}
+}

src/core/webview/webviewMessageHandler.ts

@@ -16,8 +16,10 @@ import { CloudService } from "@roo-code/cloud"
 import { TelemetryService } from "@roo-code/telemetry"
 
 import { type ApiMessage } from "../task-persistence/apiMessages"
+import { saveTaskMessages } from "../task-persistence"
 
 import { ClineProvider } from "./ClineProvider"
+import { handleCheckpointRestoreOperation } from "./checkpointRestoreHandler"
 import { changeLanguage, t } from "../../i18n"
 import { Package } from "../../shared/package"
 import { RouterName, toRouterName, ModelRecord } from "../../shared/api"
@@ -71,10 +73,10 @@ export const webviewMessageHandler = async (
 	 * Shared utility to find message indices based on timestamp
 	 */
 	const findMessageIndices = (messageTs: number, currentCline: any) => {
-		const timeCutoff = messageTs - 1000 // 1 second buffer before the message
-		const messageIndex = currentCline.clineMessages.findIndex((msg: ClineMessage) => msg.ts && msg.ts >= timeCutoff)
+		// Find the exact message by timestamp, not the first one after a cutoff
+		const messageIndex = currentCline.clineMessages.findIndex((msg: ClineMessage) => msg.ts === messageTs)
 		const apiConversationHistoryIndex = currentCline.apiConversationHistory.findIndex(
-			(msg: ApiMessage) => msg.ts && msg.ts >= timeCutoff,
+			(msg: ApiMessage) => msg.ts === messageTs,
 		)
 		return { messageIndex, apiConversationHistoryIndex }
 	}
@@ -101,38 +103,110 @@ export const webviewMessageHandler = async (
 	 * Handles message deletion operations with user confirmation
 	 */
 	const handleDeleteOperation = async (messageTs: number): Promise<void> => {
+		// Check if there's a checkpoint before this message
+		const currentCline = provider.getCurrentTask()
+		let hasCheckpoint = false
+		if (currentCline) {
+			const { messageIndex } = findMessageIndices(messageTs, currentCline)
+			if (messageIndex !== -1) {
+				// Find the last checkpoint before this message
+				const checkpoints = currentCline.clineMessages.filter(
+					(msg) => msg.say === "checkpoint_saved" && msg.ts > messageTs,
+				)
+
+				hasCheckpoint = checkpoints.length > 0
+			} 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,
 		})
 	}
 
 	/**
 	 * Handles confirmed message deletion from webview dialog
 	 */
-	const handleDeleteMessageConfirm = async (messageTs: number): Promise<void> => {
-		// Only proceed if we have a current task.
-		if (provider.getCurrentTask()) {
-			const currentCline = provider.getCurrentTask()!
-			const { messageIndex, apiConversationHistoryIndex } = findMessageIndices(messageTs, currentCline)
+	const handleDeleteMessageConfirm = async (messageTs: number, restoreCheckpoint?: boolean): Promise<void> => {
+		const currentCline = provider.getCurrentTask()
+		if (!currentCline) {
+			console.error("[handleDeleteMessageConfirm] No current cline available")
+			return
+		}
 
-			if (messageIndex !== -1) {
-				try {
-					const { historyItem } = await provider.getTaskWithId(currentCline.taskId)
+		const { messageIndex, apiConversationHistoryIndex } = findMessageIndices(messageTs, currentCline)
 
-					// Delete this message and all subsequent messages
-					await removeMessagesThisAndSubsequent(currentCline, messageIndex, apiConversationHistoryIndex)
+		if (messageIndex === -1) {
+			const errorMessage = `Message with timestamp ${messageTs} not found`
+			console.error("[handleDeleteMessageConfirm]", errorMessage)
+			await vscode.window.showErrorMessage(errorMessage)
+			return
+		}
 
-					// Initialize with history item after deletion
-					await provider.createTaskWithHistoryItem(historyItem)
-				} catch (error) {
-					console.error("Error in delete message:", error)
-					vscode.window.showErrorMessage(
-						`Error deleting message: ${error instanceof Error ? error.message : String(error)}`,
-					)
+		try {
+			const targetMessage = currentCline.clineMessages[messageIndex]
+
+			// If checkpoint restoration is requested, find and restore to the last checkpoint before this message
+			if (restoreCheckpoint) {
+				// Find the last checkpoint before this message
+				const checkpoints = currentCline.clineMessages.filter(
+					(msg) => msg.say === "checkpoint_saved" && msg.ts > messageTs,
+				)
+
+				const nextCheckpoint = checkpoints[0]
+
+				if (nextCheckpoint && nextCheckpoint.text) {
+					await handleCheckpointRestoreOperation({
+						provider,
+						currentCline,
+						messageTs: targetMessage.ts!,
+						messageIndex,
+						checkpoint: { hash: nextCheckpoint.text },
+						operation: "delete",
+					})
+				} else {
+					// No checkpoint found before this message
+					console.log("[handleDeleteMessageConfirm] No checkpoint found before message")
+					vscode.window.showWarningMessage("No checkpoint found before this message")
+				}
+			} else {
+				// For non-checkpoint deletes, preserve checkpoint associations for remaining messages
+				// Store checkpoints from messages that will be preserved
+				const preservedCheckpoints = new Map<number, any>()
+				for (let i = 0; i < messageIndex; i++) {
+					const msg = currentCline.clineMessages[i]
+					if (msg?.checkpoint && msg.ts) {
+						preservedCheckpoints.set(msg.ts, msg.checkpoint)
+					}
+				}
+
+				// Delete this message and all subsequent messages
+				await removeMessagesThisAndSubsequent(currentCline, messageIndex, apiConversationHistoryIndex)
+
+				// Restore checkpoint associations for preserved messages
+				for (const [ts, checkpoint] of preservedCheckpoints) {
+					const msgIndex = currentCline.clineMessages.findIndex((msg) => msg.ts === ts)
+					if (msgIndex !== -1) {
+						currentCline.clineMessages[msgIndex].checkpoint = checkpoint
+					}
 				}
+
+				// Save the updated messages with restored checkpoints
+				await saveTaskMessages({
+					messages: currentCline.clineMessages,
+					taskId: currentCline.taskId,
+					globalStoragePath: provider.contextProxy.globalStorageUri.fsPath,
+				})
 			}
+		} catch (error) {
+			console.error("Error in delete message:", error)
+			vscode.window.showErrorMessage(
+				`Error deleting message: ${error instanceof Error ? error.message : String(error)}`,
+			)
 		}
 	}
 
@@ -140,11 +214,31 @@ export const webviewMessageHandler = async (
 	 * Handles message editing operations with user confirmation
 	 */
 	const handleEditOperation = async (messageTs: number, editedContent: string, images?: string[]): Promise<void> => {
+		// Check if there's a checkpoint before this message
+		const currentCline = provider.getCurrentTask()
+		let hasCheckpoint = false
+		if (currentCline) {
+			const { messageIndex } = findMessageIndices(messageTs, currentCline)
+			if (messageIndex !== -1) {
+				// Find the last checkpoint before this message
+				const checkpoints = currentCline.clineMessages.filter(
+					(msg) => msg.say === "checkpoint_saved" && msg.ts > messageTs,
+				)
+
+				hasCheckpoint = checkpoints.length > 0
+			} else {
+				console.log("[webviewMessageHandler] Edit - Message not found in clineMessages!")
+			}
+		} else {
+			console.log("[webviewMessageHandler] Edit - No currentCline available!")
+		}
+
 		// Send message to webview to show edit confirmation dialog
 		await provider.postMessageToWebview({
 			type: "showEditMessageDialog",
 			messageTs,
 			text: editedContent,
+			hasCheckpoint,
 			images,
 		})
 	}
@@ -155,38 +249,105 @@ export const webviewMessageHandler = async (
 	const handleEditMessageConfirm = async (
 		messageTs: number,
 		editedContent: string,
+		restoreCheckpoint?: boolean,
 		images?: string[],
 	): Promise<void> => {
-		// Only proceed if we have a current task.
-		if (provider.getCurrentTask()) {
-			const currentCline = provider.getCurrentTask()!
+		const currentCline = provider.getCurrentTask()
+		if (!currentCline) {
+			console.error("[handleEditMessageConfirm] No current cline available")
+			return
+		}
 
-			// Use findMessageIndices to find messages based on timestamp
-			const { messageIndex, apiConversationHistoryIndex } = findMessageIndices(messageTs, currentCline)
+		// Use findMessageIndices to find messages based on timestamp
+		const { messageIndex, apiConversationHistoryIndex } = findMessageIndices(messageTs, currentCline)
 
-			if (messageIndex !== -1) {
-				try {
-					// Edit this message and delete subsequent
-					await removeMessagesThisAndSubsequent(currentCline, messageIndex, apiConversationHistoryIndex)
-
-					// Process the edited message as a regular user message
-					// This will add it to the conversation and trigger an AI response
-					webviewMessageHandler(provider, {
-						type: "askResponse",
-						askResponse: "messageResponse",
-						text: editedContent,
-						images,
+		if (messageIndex === -1) {
+			const errorMessage = `Message with timestamp ${messageTs} not found`
+			console.error("[handleEditMessageConfirm]", errorMessage)
+			await vscode.window.showErrorMessage(errorMessage)
+			return
+		}
+
+		try {
+			const targetMessage = currentCline.clineMessages[messageIndex]
+
+			// If checkpoint restoration is requested, find and restore to the last checkpoint before this message
+			if (restoreCheckpoint) {
+				// Find the last checkpoint before this message
+				const checkpoints = currentCline.clineMessages.filter(
+					(msg) => msg.say === "checkpoint_saved" && msg.ts > messageTs,
+				)
+
+				const nextCheckpoint = checkpoints[0]
+
+				if (nextCheckpoint && nextCheckpoint.text) {
+					await handleCheckpointRestoreOperation({
+						provider,
+						currentCline,
+						messageTs: targetMessage.ts!,
+						messageIndex,
+						checkpoint: { hash: nextCheckpoint.text },
+						operation: "edit",
+						editData: {
+							editedContent,
+							images,
+							apiConversationHistoryIndex,
+						},
 					})
+					// The task will be cancelled and reinitialized by checkpointRestore
+					// The pending edit will be processed in the reinitialized task
+					return
+				} else {
+					// No checkpoint found before this message
+					console.log("[handleEditMessageConfirm] No checkpoint found before message")
+					vscode.window.showWarningMessage("No checkpoint found before this message")
+					// Continue with non-checkpoint edit
+				}
+			}
 
-					// Don't initialize with history item for edit operations
-					// The webviewMessageHandler will handle the conversation state
-				} catch (error) {
-					console.error("Error in edit message:", error)
-					vscode.window.showErrorMessage(
-						`Error editing message: ${error instanceof Error ? error.message : String(error)}`,
-					)
+			// For non-checkpoint edits, preserve checkpoint associations for remaining messages
+			// Store checkpoints from messages that will be preserved
+			const preservedCheckpoints = new Map<number, any>()
+			for (let i = 0; i < messageIndex; i++) {
+				const msg = currentCline.clineMessages[i]
+				if (msg?.checkpoint && msg.ts) {
+					preservedCheckpoints.set(msg.ts, msg.checkpoint)
+				}
+			}
+
+			// Edit this message and delete subsequent
+			await removeMessagesThisAndSubsequent(currentCline, messageIndex, apiConversationHistoryIndex)
+
+			// Restore checkpoint associations for preserved messages
+			for (const [ts, checkpoint] of preservedCheckpoints) {
+				const msgIndex = currentCline.clineMessages.findIndex((msg) => msg.ts === ts)
+				if (msgIndex !== -1) {
+					currentCline.clineMessages[msgIndex].checkpoint = checkpoint
 				}
 			}
+
+			// Save the updated messages with restored checkpoints
+			await saveTaskMessages({
+				messages: currentCline.clineMessages,
+				taskId: currentCline.taskId,
+				globalStoragePath: provider.contextProxy.globalStorageUri.fsPath,
+			})
+
+			// Process the edited message as a regular user message
+			webviewMessageHandler(provider, {
+				type: "askResponse",
+				askResponse: "messageResponse",
+				text: editedContent,
+				images,
+			})
+
+			// Don't initialize with history item for edit operations
+			// The webviewMessageHandler will handle the conversation state
+		} catch (error) {
+			console.error("Error in edit message:", error)
+			vscode.window.showErrorMessage(
+				`Error editing message: ${error instanceof Error ? error.message : String(error)}`,
+			)
 		}
 	}
 
@@ -1674,12 +1835,17 @@ 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, message.images)
+				await handleEditMessageConfirm(
+					message.messageTs,
+					message.text,
+					message.restoreCheckpoint,
+					message.images,
+				)
 			}
 			break
 		case "getListApiConfiguration":

src/shared/ExtensionMessage.ts

@@ -195,6 +195,7 @@ export interface ExtensionMessage {
 	rulesFolderPath?: string
 	settings?: any
 	messageTs?: number
+	hasCheckpoint?: boolean
 	context?: string
 	commands?: Command[]
 	queuedMessages?: QueuedMessage[]

src/shared/WebviewMessage.ts

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

webview-ui/src/App.tsx

@@ -19,6 +19,7 @@ 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 { CheckpointRestoreDialog } from "./components/chat/CheckpointRestoreDialog"
 import { DeleteMessageDialog, EditMessageDialog } from "./components/chat/MessageModificationConfirmationDialog"
 import ErrorBoundary from "./components/ErrorBoundary"
 import { CloudView } from "./components/cloud/CloudView"
@@ -37,18 +38,21 @@ interface HumanRelayDialogState {
 interface DeleteMessageDialogState {
 	isOpen: boolean
 	messageTs: number
+	hasCheckpoint: boolean
 }
 
 interface EditMessageDialogState {
 	isOpen: boolean
 	messageTs: number
 	text: string
+	hasCheckpoint: boolean
 	images?: string[]
 }
 
 // Memoize dialog components to prevent unnecessary re-renders
 const MemoizedDeleteMessageDialog = React.memo(DeleteMessageDialog)
 const MemoizedEditMessageDialog = React.memo(EditMessageDialog)
+const MemoizedCheckpointRestoreDialog = React.memo(CheckpointRestoreDialog)
 const MemoizedHumanRelayDialog = React.memo(HumanRelayDialog)
 
 const tabsByMessageAction: Partial<Record<NonNullable<ExtensionMessage["action"]>, Tab>> = {
@@ -91,12 +95,14 @@ const App = () => {
 	const [deleteMessageDialogState, setDeleteMessageDialogState] = useState<DeleteMessageDialogState>({
 		isOpen: false,
 		messageTs: 0,
+		hasCheckpoint: false,
 	})
 
 	const [editMessageDialogState, setEditMessageDialogState] = useState<EditMessageDialogState>({
 		isOpen: false,
 		messageTs: 0,
 		text: "",
+		hasCheckpoint: false,
 		images: [],
 	})
 
@@ -159,7 +165,11 @@ const App = () => {
 			}
 
 			if (message.type === "showDeleteMessageDialog" && message.messageTs) {
-				setDeleteMessageDialogState({ isOpen: true, messageTs: message.messageTs })
+				setDeleteMessageDialogState({
+					isOpen: true,
+					messageTs: message.messageTs,
+					hasCheckpoint: message.hasCheckpoint || false,
+				})
 			}
 
 			if (message.type === "showEditMessageDialog" && message.messageTs && message.text) {
@@ -167,6 +177,7 @@ const App = () => {
 					isOpen: true,
 					messageTs: message.messageTs,
 					text: message.text,
+					hasCheckpoint: message.hasCheckpoint || false,
 					images: message.images || [],
 				})
 			}
@@ -271,30 +282,65 @@ const App = () => {
 				onSubmit={(requestId, text) => vscode.postMessage({ type: "humanRelayResponse", requestId, text })}
 				onCancel={(requestId) => vscode.postMessage({ type: "humanRelayCancel", requestId })}
 			/>
-			<MemoizedDeleteMessageDialog
-				open={deleteMessageDialogState.isOpen}
-				onOpenChange={(open) => setDeleteMessageDialogState((prev) => ({ ...prev, isOpen: open }))}
-				onConfirm={() => {
-					vscode.postMessage({
-						type: "deleteMessageConfirm",
-						messageTs: deleteMessageDialogState.messageTs,
-					})
-					setDeleteMessageDialogState((prev) => ({ ...prev, isOpen: false }))
-				}}
-			/>
-			<MemoizedEditMessageDialog
-				open={editMessageDialogState.isOpen}
-				onOpenChange={(open) => setEditMessageDialogState((prev) => ({ ...prev, isOpen: open }))}
-				onConfirm={() => {
-					vscode.postMessage({
-						type: "editMessageConfirm",
-						messageTs: editMessageDialogState.messageTs,
-						text: editMessageDialogState.text,
-						images: editMessageDialogState.images,
-					})
-					setEditMessageDialogState((prev) => ({ ...prev, isOpen: false }))
-				}}
-			/>
+			{deleteMessageDialogState.hasCheckpoint ? (
+				<MemoizedCheckpointRestoreDialog
+					open={deleteMessageDialogState.isOpen}
+					type="delete"
+					hasCheckpoint={deleteMessageDialogState.hasCheckpoint}
+					onOpenChange={(open: boolean) => setDeleteMessageDialogState((prev) => ({ ...prev, isOpen: open }))}
+					onConfirm={(restoreCheckpoint: boolean) => {
+						vscode.postMessage({
+							type: "deleteMessageConfirm",
+							messageTs: deleteMessageDialogState.messageTs,
+							restoreCheckpoint,
+						})
+						setDeleteMessageDialogState((prev) => ({ ...prev, isOpen: false }))
+					}}
+				/>
+			) : (
+				<MemoizedDeleteMessageDialog
+					open={deleteMessageDialogState.isOpen}
+					onOpenChange={(open: boolean) => setDeleteMessageDialogState((prev) => ({ ...prev, isOpen: open }))}
+					onConfirm={() => {
+						vscode.postMessage({
+							type: "deleteMessageConfirm",
+							messageTs: deleteMessageDialogState.messageTs,
+						})
+						setDeleteMessageDialogState((prev) => ({ ...prev, isOpen: false }))
+					}}
+				/>
+			)}
+			{editMessageDialogState.hasCheckpoint ? (
+				<MemoizedCheckpointRestoreDialog
+					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 }))
+					}}
+				/>
+			) : (
+				<MemoizedEditMessageDialog
+					open={editMessageDialogState.isOpen}
+					onOpenChange={(open: boolean) => setEditMessageDialogState((prev) => ({ ...prev, isOpen: open }))}
+					onConfirm={() => {
+						vscode.postMessage({
+							type: "editMessageConfirm",
+							messageTs: editMessageDialogState.messageTs,
+							text: editMessageDialogState.text,
+							images: editMessageDialogState.images,
+						})
+						setEditMessageDialogState((prev) => ({ ...prev, isOpen: false }))
+					}}
+				/>
+			)}
 		</>
 	)
 }

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

@@ -5,6 +5,7 @@ import deepEqual from "fast-deep-equal"
 import { VSCodeBadge, VSCodeButton } from "@vscode/webview-ui-toolkit/react"
 
 import type { ClineMessage, FollowUpData, SuggestionItem } from "@roo-code/types"
+import { Mode } from "@roo/modes"
 
 import { ClineApiReqInfo, ClineAskUseMcpServer, ClineSayTool } from "@roo/ExtensionMessage"
 import { COMMAND_OUTPUT_STRING } from "@roo/combineCommandSequences"
@@ -41,7 +42,11 @@ import { CommandExecutionError } from "./CommandExecutionError"
 import { AutoApprovedRequestLimitWarning } from "./AutoApprovedRequestLimitWarning"
 import { CondenseContextErrorRow, CondensingContextRow, ContextCondenseRow } from "./ContextCondenseRow"
 import CodebaseSearchResultsDisplay from "./CodebaseSearchResultsDisplay"
+import { appendImages } from "@src/utils/imageUtils"
 import { McpExecution } from "./McpExecution"
+import { ChatTextArea } from "./ChatTextArea"
+import { MAX_IMAGES_PER_MESSAGE } from "./ChatView"
+import { useSelectedModel } from "../ui/hooks/useSelectedModel"
 
 interface ChatRowProps {
 	message: ClineMessage
@@ -111,19 +116,70 @@ export const ChatRowContent = ({
 }: ChatRowContentProps) => {
 	const { t } = useTranslation()
 
-	const { mcpServers, alwaysAllowMcp, currentCheckpoint } = useExtensionState()
-
+	const { mcpServers, alwaysAllowMcp, currentCheckpoint, mode, apiConfiguration } = useExtensionState()
+	const { info: model } = useSelectedModel(apiConfiguration)
 	const [reasoningCollapsed, setReasoningCollapsed] = useState(true)
 	const [isDiffErrorExpanded, setIsDiffErrorExpanded] = useState(false)
 	const [showCopySuccess, setShowCopySuccess] = useState(false)
-
+	const [isEditing, setIsEditing] = useState(false)
+	const [editedContent, setEditedContent] = useState("")
+	const [editMode, setEditMode] = useState<Mode>(mode || "code")
+	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)
 	}, [onToggleExpand, message.ts])
 
+	// Handle edit button click
+	const handleEditClick = useCallback(() => {
+		setIsEditing(true)
+		setEditedContent(message.text || "")
+		setEditImages(message.images || [])
+		setEditMode(mode || "code")
+		// Edit mode is now handled entirely in the frontend
+		// No need to notify the backend
+	}, [message.text, message.images, mode])
+
+	// Handle cancel edit
+	const handleCancelEdit = useCallback(() => {
+		setIsEditing(false)
+		setEditedContent(message.text || "")
+		setEditImages(message.images || [])
+		setEditMode(mode || "code")
+	}, [message.text, message.images, mode])
+
+	// Handle save edit
+	const handleSaveEdit = useCallback(() => {
+		setIsEditing(false)
+		// Send edited message to backend
+		vscode.postMessage({
+			type: "submitEditedMessage",
+			value: message.ts,
+			editedMessageContent: editedContent,
+			images: editImages,
+		})
+	}, [message.ts, editedContent, editImages])
+
+	// Handle image selection for editing
+	const handleSelectImages = useCallback(() => {
+		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") {
 			const info = safeJsonParse<ClineApiReqInfo>(message.text)
@@ -1117,26 +1173,58 @@ export const ChatRowContent = ({
 				case "user_feedback":
 					return (
 						<div className="bg-vscode-editor-background border rounded-xs p-1 overflow-hidden whitespace-pre-wrap">
-							<div className="flex justify-between">
-								<div className="flex-grow px-2 py-1 wrap-anywhere">
-									<Mention text={message.text} withShadow />
+							{isEditing ? (
+								<div className="flex flex-col gap-2 p-2">
+									<ChatTextArea
+										inputValue={editedContent}
+										setInputValue={setEditedContent}
+										sendingDisabled={false}
+										selectApiConfigDisabled={true}
+										placeholderText={t("chat:editMessage.placeholder")}
+										selectedImages={editImages}
+										setSelectedImages={setEditImages}
+										onSend={handleSaveEdit}
+										onSelectImages={handleSelectImages}
+										shouldDisableImages={!model?.supportsImages}
+										mode={editMode}
+										setMode={setEditMode}
+										modeShortcutText=""
+										isEditMode={true}
+										onCancel={handleCancelEdit}
+									/>
 								</div>
-								<div className="flex">
-									<Button
-										variant="ghost"
-										size="icon"
-										className="shrink-0"
-										disabled={isStreaming}
-										onClick={(e) => {
-											e.stopPropagation()
-											vscode.postMessage({ type: "deleteMessage", value: message.ts })
-										}}>
-										<span className="codicon codicon-trash" />
-									</Button>
+							) : (
+								<div className="flex justify-between">
+									<div className="flex-grow px-2 py-1 wrap-anywhere">
+										<Mention text={message.text} withShadow />
+									</div>
+									<div className="flex">
+										<Button
+											variant="ghost"
+											size="icon"
+											className="shrink-0"
+											disabled={isStreaming}
+											onClick={(e) => {
+												e.stopPropagation()
+												handleEditClick()
+											}}>
+											<span className="codicon codicon-edit" />
+										</Button>
+										<Button
+											variant="ghost"
+											size="icon"
+											className="shrink-0"
+											disabled={isStreaming}
+											onClick={(e) => {
+												e.stopPropagation()
+												vscode.postMessage({ type: "deleteMessage", value: message.ts })
+											}}>
+											<span className="codicon codicon-trash" />
+										</Button>
+									</div>
 								</div>
-							</div>
-
-							{message.images && message.images.length > 0 && (
+							)}
+							{!isEditing && message.images && message.images.length > 0 && (
 								<Thumbnails images={message.images} style={{ marginTop: "8px" }} />
 							)}
 						</div>

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

@@ -1,7 +1,7 @@
 import React, { forwardRef, useCallback, useEffect, useLayoutEffect, useMemo, useRef, useState } from "react"
 import { useEvent } from "react-use"
 import DynamicTextArea from "react-textarea-autosize"
-import { VolumeX, Image, WandSparkles, SendHorizontal } from "lucide-react"
+import { VolumeX, Image, WandSparkles, SendHorizontal, MessageSquareX } from "lucide-react"
 
 import { mentionRegex, mentionRegexGlobal, commandRegexGlobal, unescapeSpaces } from "@roo/context-mentions"
 import { WebviewMessage } from "@roo/WebviewMessage"
@@ -47,6 +47,9 @@ interface ChatTextAreaProps {
 	mode: Mode
 	setMode: (value: Mode) => void
 	modeShortcutText: string
+	// Edit mode props
+	isEditMode?: boolean
+	onCancel?: () => void
 }
 
 export const ChatTextArea = forwardRef<HTMLTextAreaElement, ChatTextAreaProps>(
@@ -65,6 +68,8 @@ export const ChatTextArea = forwardRef<HTMLTextAreaElement, ChatTextAreaProps>(
 			mode,
 			setMode,
 			modeShortcutText,
+			isEditMode = false,
+			onCancel,
 		},
 		ref,
 	) => {
@@ -881,6 +886,7 @@ export const ChatTextArea = forwardRef<HTMLTextAreaElement, ChatTextAreaProps>(
 
 		const placeholderBottomText = `\n(${t("chat:addContext")}${shouldDisableImages ? `, ${t("chat:dragFiles")}` : `, ${t("chat:dragFilesImages")}`})`
 
+		// Common mode selector handler
 		const handleModeChange = useCallback(
 			(value: Mode) => {
 				setMode(value)
@@ -889,6 +895,7 @@ export const ChatTextArea = forwardRef<HTMLTextAreaElement, ChatTextAreaProps>(
 			[setMode],
 		)
 
+		// Helper function to handle API config change
 		const handleApiConfigChange = useCallback((value: string) => {
 			vscode.postMessage({ type: "loadApiConfigurationById", text: value })
 		}, [])
@@ -1096,6 +1103,27 @@ export const ChatTextArea = forwardRef<HTMLTextAreaElement, ChatTextAreaProps>(
 							</div>
 
 							<div className="absolute bottom-1 right-1 z-30">
+								{isEditMode && (
+									<StandardTooltip content={t("chat:cancel.title")}>
+										<button
+											aria-label={t("chat:cancel.title")}
+											disabled={false}
+											onClick={onCancel}
+											className={cn(
+												"relative inline-flex items-center justify-center",
+												"bg-transparent border-none p-1.5",
+												"rounded-md min-w-[28px] min-h-[28px]",
+												"opacity-60 hover:opacity-100 text-vscode-descriptionForeground hover:text-vscode-foreground",
+												"transition-all duration-150",
+												"hover:bg-[rgba(255,255,255,0.03)] hover:border-[rgba(255,255,255,0.15)]",
+												"focus:outline-none focus-visible:ring-1 focus-visible:ring-vscode-focusBorder",
+												"active:bg-[rgba(255,255,255,0.1)]",
+												"cursor-pointer",
+											)}>
+											<MessageSquareX className="w-4 h-4" />
+										</button>
+									</StandardTooltip>
+								)}
 								<StandardTooltip content={t("chat:sendMessage")}>
 									<button
 										aria-label={t("chat:sendMessage")}
@@ -1193,8 +1221,8 @@ export const ChatTextArea = forwardRef<HTMLTextAreaElement, ChatTextAreaProps>(
 								</button>
 							</StandardTooltip>
 						)}
-						<SlashCommandsPopover />
-						<IndexingStatusBadge />
+						{!isEditMode ? <SlashCommandsPopover /> : null}
+						{!isEditMode ? <IndexingStatusBadge /> : null}
 						<StandardTooltip content={t("chat:addImages")}>
 							<button
 								aria-label={t("chat:addImages")}

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

@@ -843,9 +843,30 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
 	useMount(() => textAreaRef.current?.focus())
 
 	const visibleMessages = useMemo(() => {
+		// Pre-compute checkpoint hashes that have associated user messages for O(1) lookup
+		const userMessageCheckpointHashes = new Set<string>()
+		modifiedMessages.forEach((msg) => {
+			if (
+				msg.say === "user_feedback" &&
+				msg.checkpoint &&
+				(msg.checkpoint as any).type === "user_message" &&
+				(msg.checkpoint as any).hash
+			) {
+				userMessageCheckpointHashes.add((msg.checkpoint as any).hash)
+			}
+		})
+
 		// Remove the 500-message limit to prevent array index shifting
 		// Virtuoso is designed to efficiently handle large lists through virtualization
-		const newVisibleMessages = modifiedMessages.filter((message: ClineMessage) => {
+		const newVisibleMessages = modifiedMessages.filter((message) => {
+			// Filter out checkpoint_saved messages that are associated with user messages
+			if (message.say === "checkpoint_saved" && message.text) {
+				// Use O(1) Set lookup instead of O(n) array search
+				if (userMessageCheckpointHashes.has(message.text)) {
+					return false
+				}
+			}
+
 			if (everVisibleMessagesTsRef.current.has(message.ts)) {
 				const alwaysHiddenOnceProcessedAsk: ClineAsk[] = [
 					"api_req_failed",

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.editMessage") : t("common:confirmation.deleteMessage")
+	const description = isEdit
+		? t("common:confirmation.editQuestionWithCheckpoint")
+		: t("common:confirmation.deleteQuestionWithCheckpoint")
+
+	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.editOnly") : t("common:confirmation.deleteOnly")}
+					</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.restoreToCheckpoint")}
+						</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,245 @@
+// 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.deleteMessage": "Delete Message",
+				"common:confirmation.editMessage": "Edit Message",
+				"common:confirmation.deleteQuestionWithCheckpoint":
+					"Deleting this message will delete all subsequent messages in the conversation. Do you want to proceed?",
+				"common:confirmation.editQuestionWithCheckpoint":
+					"Editing this message will delete all subsequent messages in the conversation. Do you want to proceed?",
+				"common:confirmation.editOnly": "Edit Only",
+				"common:confirmation.deleteOnly": "Delete Only",
+				"common:confirmation.restoreToCheckpoint": "Restore to Checkpoint",
+				"common:answers.cancel": "Cancel",
+			}
+			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("Edit Only")).toBeInTheDocument()
+			expect(screen.getByText("Cancel")).toBeInTheDocument()
+			expect(screen.queryByText("Restore to 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("Delete Only")).toBeInTheDocument()
+			expect(screen.getByText("Cancel")).toBeInTheDocument()
+			expect(screen.queryByText("Restore to 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("Edit Only")).toBeInTheDocument()
+			expect(screen.getByText("Restore to Checkpoint")).toBeInTheDocument()
+			expect(screen.getByText("Cancel")).toBeInTheDocument()
+		})
+
+		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("Delete Only")).toBeInTheDocument()
+			expect(screen.getByText("Restore to Checkpoint")).toBeInTheDocument()
+			expect(screen.getByText("Cancel")).toBeInTheDocument()
+		})
+	})
+
+	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 edit only is clicked", () => {
+			const onConfirm = vi.fn()
+			render(<CheckpointRestoreDialog {...defaultProps} onConfirm={onConfirm} />)
+
+			fireEvent.click(screen.getByText("Edit Only"))
+			expect(onConfirm).toHaveBeenCalledWith(false) // restoreCheckpoint
+		})
+
+		it("calls onConfirm with restoreCheckpoint=false when edit only is clicked with checkpoint", () => {
+			const onConfirm = vi.fn()
+			render(<CheckpointRestoreDialog {...defaultProps} onConfirm={onConfirm} hasCheckpoint={true} />)
+
+			fireEvent.click(screen.getByText("Edit Only"))
+			expect(onConfirm).toHaveBeenCalledWith(false) // restoreCheckpoint
+		})
+
+		it("calls onConfirm with restoreCheckpoint=true when restore to checkpoint is clicked", () => {
+			const onConfirm = vi.fn()
+			render(<CheckpointRestoreDialog {...defaultProps} onConfirm={onConfirm} hasCheckpoint={true} />)
+
+			fireEvent.click(screen.getByText("Restore to Checkpoint"))
+			expect(onConfirm).toHaveBeenCalledWith(true) // restoreCheckpoint
+		})
+
+		it("calls onOpenChange when dialog is closed", () => {
+			const onOpenChange = vi.fn()
+			render(<CheckpointRestoreDialog {...defaultProps} onOpenChange={onOpenChange} hasCheckpoint={true} />)
+
+			fireEvent.click(screen.getByText("Edit Only"))
+			expect(onOpenChange).toHaveBeenCalledWith(false)
+		})
+	})
+
+	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("maintains state when dialog stays open", async () => {
+			const { rerender } = render(<CheckpointRestoreDialog {...defaultProps} hasCheckpoint={true} open={true} />)
+
+			// Verify initial state
+			expect(screen.getByText("Edit Only")).toBeInTheDocument()
+			expect(screen.getByText("Restore to Checkpoint")).toBeInTheDocument()
+
+			// Re-render with same props
+			rerender(<CheckpointRestoreDialog {...defaultProps} hasCheckpoint={true} open={true} />)
+
+			// Should still have same buttons
+			expect(screen.getByText("Edit Only")).toBeInTheDocument()
+			expect(screen.getByText("Restore to Checkpoint")).toBeInTheDocument()
+		})
+	})
+
+	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.getByRole("button", { name: "Edit Only" })).toBeInTheDocument()
+			expect(screen.getByRole("button", { name: "Restore to Checkpoint" })).toBeInTheDocument()
+			expect(screen.getByRole("button", { name: "Cancel" })).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 button clicks", async () => {
+			const onConfirm = vi.fn()
+			const onOpenChange = vi.fn()
+			render(
+				<CheckpointRestoreDialog
+					{...defaultProps}
+					onConfirm={onConfirm}
+					onOpenChange={onOpenChange}
+					hasCheckpoint={true}
+				/>,
+			)
+
+			const editOnlyButton = screen.getByText("Edit Only")
+
+			// Click button once
+			fireEvent.click(editOnlyButton)
+
+			// Should be called once with correct parameters
+			expect(onConfirm).toHaveBeenCalledTimes(1)
+			expect(onConfirm).toHaveBeenCalledWith(false) // restoreCheckpoint
+			expect(onOpenChange).toHaveBeenCalledWith(false) // dialog should close
+		})
+	})
+
+	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/ca/common.json

@@ -72,7 +72,13 @@
 		"deleteWarning": "Eliminar aquest missatge eliminarà tots els missatges posteriors de la conversa. Vols continuar?",
 		"editMessage": "Editar missatge",
 		"editWarning": "Editar aquest missatge eliminarà tots els missatges posteriors de la conversa. Vols continuar?",
-		"proceed": "Continuar"
+		"editQuestionWithCheckpoint": "Editar aquest missatge eliminarà tots els missatges posteriors de la conversa. També vols desfer tots els canvis fins a aquest punt de control?",
+		"deleteQuestionWithCheckpoint": "Eliminar aquest missatge eliminarà tots els missatges posteriors de la conversa. També vols desfer tots els canvis fins a aquest punt de control?",
+		"editOnly": "No, només editar el missatge",
+		"deleteOnly": "No, només eliminar el missatge",
+		"restoreToCheckpoint": "Sí, restaurar el punt de control",
+		"proceed": "Continuar",
+		"dontShowAgain": "No tornis a mostrar això"
 	},
 	"time_ago": {
 		"just_now": "ara mateix",

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

@@ -72,7 +72,13 @@
 		"deleteWarning": "Das Löschen dieser Nachricht wird alle nachfolgenden Nachrichten in der Unterhaltung löschen. Möchtest du fortfahren?",
 		"editMessage": "Nachricht bearbeiten",
 		"editWarning": "Das Bearbeiten dieser Nachricht wird alle nachfolgenden Nachrichten in der Unterhaltung löschen. Möchtest du fortfahren?",
-		"proceed": "Fortfahren"
+		"editQuestionWithCheckpoint": "Das Bearbeiten dieser Nachricht wird alle späteren Nachrichten in der Unterhaltung löschen. Möchtest du auch alle Änderungen bis zu diesem Checkpoint rückgängig machen?",
+		"deleteQuestionWithCheckpoint": "Das Löschen dieser Nachricht wird alle späteren Nachrichten in der Unterhaltung löschen. Möchtest du auch alle Änderungen bis zu diesem Checkpoint rückgängig machen?",
+		"editOnly": "Nein, nur Nachricht bearbeiten",
+		"deleteOnly": "Nein, nur Nachricht löschen",
+		"restoreToCheckpoint": "Ja, Checkpoint wiederherstellen",
+		"proceed": "Fortfahren",
+		"dontShowAgain": "Nicht mehr anzeigen"
 	},
 	"time_ago": {
 		"just_now": "gerade eben",

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

@@ -72,7 +72,13 @@
 		"deleteWarning": "Deleting this message will delete all subsequent messages in the conversation. Do you want to proceed?",
 		"editMessage": "Edit Message",
 		"editWarning": "Editing this message will delete all subsequent messages in the conversation. Do you want to proceed?",
-		"proceed": "Proceed"
+		"editQuestionWithCheckpoint": "Editing this message will delete all later messages in the conversation. Do you also want to undo all changes back to this checkpoint?",
+		"deleteQuestionWithCheckpoint": "Deleting this message will delete all later messages in the conversation. Do you also want to undo all changes back to this checkpoint?",
+		"editOnly": "No, edit message only",
+		"deleteOnly": "No, delete message only",
+		"restoreToCheckpoint": "Yes, restore the checkpoint",
+		"proceed": "Proceed",
+		"dontShowAgain": "Don't show this again"
 	},
 	"time_ago": {
 		"just_now": "just now",

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

@@ -72,7 +72,13 @@
 		"deleteWarning": "Eliminar este mensaje eliminará todos los mensajes posteriores en la conversación. ¿Deseas continuar?",
 		"editMessage": "Editar mensaje",
 		"editWarning": "Editar este mensaje eliminará todos los mensajes posteriores en la conversación. ¿Deseas continuar?",
-		"proceed": "Continuar"
+		"editQuestionWithCheckpoint": "Editar este mensaje eliminará todos los mensajes posteriores en la conversación. ¿También deseas deshacer todos los cambios hasta este punto de control?",
+		"deleteQuestionWithCheckpoint": "Eliminar este mensaje eliminará todos los mensajes posteriores en la conversación. ¿También deseas deshacer todos los cambios hasta este punto de control?",
+		"editOnly": "No, solo editar el mensaje",
+		"deleteOnly": "No, solo eliminar el mensaje",
+		"restoreToCheckpoint": "Sí, restaurar el punto de control",
+		"proceed": "Continuar",
+		"dontShowAgain": "No mostrar esto de nuevo"
 	},
 	"time_ago": {
 		"just_now": "ahora mismo",

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

@@ -72,7 +72,13 @@
 		"deleteWarning": "Supprimer ce message supprimera tous les messages suivants dans la conversation. Voulez-vous continuer ?",
 		"editMessage": "Modifier le message",
 		"editWarning": "Modifier ce message supprimera tous les messages suivants dans la conversation. Voulez-vous continuer ?",
-		"proceed": "Continuer"
+		"editQuestionWithCheckpoint": "Modifier ce message supprimera tous les messages ultérieurs dans la conversation. Voulez-vous aussi annuler tous les changements jusqu'à ce point de contrôle ?",
+		"deleteQuestionWithCheckpoint": "Supprimer ce message supprimera tous les messages ultérieurs dans la conversation. Voulez-vous aussi annuler tous les changements jusqu'à ce point de contrôle ?",
+		"editOnly": "Non, modifier le message seulement",
+		"deleteOnly": "Non, supprimer le message seulement",
+		"restoreToCheckpoint": "Oui, restaurer le point de contrôle",
+		"proceed": "Continuer",
+		"dontShowAgain": "Ne plus afficher ceci"
 	},
 	"time_ago": {
 		"just_now": "à l'instant",

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

@@ -72,7 +72,13 @@
 		"deleteWarning": "इस संदेश को हटाने से बातचीत के सभी बाद के संदेश हट जाएंगे। क्या आप जारी रखना चाहते हैं?",
 		"editMessage": "संदेश संपादित करें",
 		"editWarning": "इस संदेश को संपादित करने से बातचीत के सभी बाद के संदेश हट जाएंगे। क्या आप जारी रखना चाहते हैं?",
-		"proceed": "जारी रखें"
+		"editQuestionWithCheckpoint": "इस संदेश को संपादित करने से बातचीत के सभी बाद के संदेश हट जाएंगे। क्या आप इस चेकपॉइंट तक सभी परिवर्तनों को भी पूर्ववत करना चाहते हैं?",
+		"deleteQuestionWithCheckpoint": "इस संदेश को हटाने से बातचीत के सभी बाद के संदेश हट जाएंगे। क्या आप इस चेकपॉइंट तक सभी परिवर्तनों को भी पूर्ववत करना चाहते हैं?",
+		"editOnly": "नहीं, केवल संदेश संपादित करें",
+		"deleteOnly": "नहीं, केवल संदेश हटाएं",
+		"restoreToCheckpoint": "हां, चेकपॉइंट पुनर्स्थापित करें",
+		"proceed": "जारी रखें",
+		"dontShowAgain": "यह फिर से न दिखाएं"
 	},
 	"time_ago": {
 		"just_now": "अभी",

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

@@ -72,7 +72,13 @@
 		"deleteWarning": "Menghapus pesan ini akan menghapus semua pesan selanjutnya dalam percakapan. Apakah kamu ingin melanjutkan?",
 		"editMessage": "Edit Pesan",
 		"editWarning": "Mengedit pesan ini akan menghapus semua pesan selanjutnya dalam percakapan. Apakah kamu ingin melanjutkan?",
-		"proceed": "Lanjutkan"
+		"editQuestionWithCheckpoint": "Mengedit pesan ini akan menghapus semua pesan selanjutnya dalam percakapan. Apakah kamu juga ingin membatalkan semua perubahan kembali ke checkpoint ini?",
+		"deleteQuestionWithCheckpoint": "Menghapus pesan ini akan menghapus semua pesan selanjutnya dalam percakapan. Apakah kamu juga ingin membatalkan semua perubahan kembali ke checkpoint ini?",
+		"editOnly": "Tidak, edit pesan saja",
+		"deleteOnly": "Tidak, hapus pesan saja",
+		"restoreToCheckpoint": "Ya, pulihkan checkpoint",
+		"proceed": "Lanjutkan",
+		"dontShowAgain": "Jangan tampilkan lagi"
 	},
 	"time_ago": {
 		"just_now": "baru saja",

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

@@ -72,7 +72,13 @@
 		"deleteWarning": "Eliminando questo messaggio verranno eliminati tutti i messaggi successivi nella conversazione. Vuoi procedere?",
 		"editMessage": "Modifica Messaggio",
 		"editWarning": "Modificando questo messaggio verranno eliminati tutti i messaggi successivi nella conversazione. Vuoi procedere?",
-		"proceed": "Procedi"
+		"editQuestionWithCheckpoint": "Modificando questo messaggio verranno eliminati tutti i messaggi successivi nella conversazione. Vuoi anche annullare tutte le modifiche fino a questo checkpoint?",
+		"deleteQuestionWithCheckpoint": "Eliminando questo messaggio verranno eliminati tutti i messaggi successivi nella conversazione. Vuoi anche annullare tutte le modifiche fino a questo checkpoint?",
+		"editOnly": "No, modifica solo il messaggio",
+		"deleteOnly": "No, elimina solo il messaggio",
+		"restoreToCheckpoint": "Sì, ripristina il checkpoint",
+		"proceed": "Procedi",
+		"dontShowAgain": "Non mostrare più"
 	},
 	"time_ago": {
 		"just_now": "proprio ora",

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

@@ -72,7 +72,13 @@
 		"deleteWarning": "このメッセージを削除すると、会話内の後続のメッセージもすべて削除されます。続行しますか？",
 		"editMessage": "メッセージを編集",
 		"editWarning": "このメッセージを編集すると、会話内の後続のメッセージもすべて削除されます。続行しますか？",
-		"proceed": "続行"
+		"editQuestionWithCheckpoint": "このメッセージを編集すると、会話内の後続のメッセージもすべて削除されます。このチェックポイントまでのすべての変更も元に戻しますか？",
+		"deleteQuestionWithCheckpoint": "このメッセージを削除すると、会話内の後続のメッセージもすべて削除されます。このチェックポイントまでのすべての変更も元に戻しますか？",
+		"editOnly": "いいえ、メッセージのみ編集",
+		"deleteOnly": "いいえ、メッセージのみ削除",
+		"restoreToCheckpoint": "はい、チェックポイントを復元",
+		"proceed": "続行",
+		"dontShowAgain": "今後表示しない"
 	},
 	"time_ago": {
 		"just_now": "たった今",

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

@@ -72,7 +72,13 @@
 		"deleteWarning": "이 메시지를 삭제하면 대화의 모든 후속 메시지가 삭제됩니다. 계속하시겠습니까?",
 		"editMessage": "메시지 편집",
 		"editWarning": "이 메시지를 편집하면 대화의 모든 후속 메시지가 삭제됩니다. 계속하시겠습니까?",
-		"proceed": "계속"
+		"editQuestionWithCheckpoint": "이 메시지를 편집하면 대화의 모든 후속 메시지가 삭제됩니다. 이 체크포인트까지의 모든 변경사항도 되돌리시겠습니까?",
+		"deleteQuestionWithCheckpoint": "이 메시지를 삭제하면 대화의 모든 후속 메시지가 삭제됩니다. 이 체크포인트까지의 모든 변경사항도 되돌리시겠습니까?",
+		"editOnly": "아니요, 메시지만 편집",
+		"deleteOnly": "아니요, 메시지만 삭제",
+		"restoreToCheckpoint": "예, 체크포인트 복원",
+		"proceed": "계속",
+		"dontShowAgain": "다시 표시하지 않음"
 	},
 	"time_ago": {
 		"just_now": "방금",

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

@@ -72,7 +72,13 @@
 		"deleteWarning": "Het verwijderen van dit bericht zal alle volgende berichten in het gesprek verwijderen. Wil je doorgaan?",
 		"editMessage": "Bericht Bewerken",
 		"editWarning": "Het bewerken van dit bericht zal alle volgende berichten in het gesprek verwijderen. Wil je doorgaan?",
-		"proceed": "Doorgaan"
+		"editQuestionWithCheckpoint": "Het bewerken van dit bericht zal alle latere berichten in het gesprek verwijderen. Wil je ook alle wijzigingen ongedaan maken tot dit checkpoint?",
+		"deleteQuestionWithCheckpoint": "Het verwijderen van dit bericht zal alle latere berichten in het gesprek verwijderen. Wil je ook alle wijzigingen ongedaan maken tot dit checkpoint?",
+		"editOnly": "Nee, alleen bericht bewerken",
+		"deleteOnly": "Nee, alleen bericht verwijderen",
+		"restoreToCheckpoint": "Ja, checkpoint herstellen",
+		"proceed": "Doorgaan",
+		"dontShowAgain": "Niet meer tonen"
 	},
 	"time_ago": {
 		"just_now": "zojuist",

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

@@ -72,7 +72,13 @@
 		"deleteWarning": "Usunięcie tej wiadomości spowoduje usunięcie wszystkich kolejnych wiadomości w rozmowie. Czy chcesz kontynuować?",
 		"editMessage": "Edytuj Wiadomość",
 		"editWarning": "Edycja tej wiadomości spowoduje usunięcie wszystkich kolejnych wiadomości w rozmowie. Czy chcesz kontynuować?",
-		"proceed": "Kontynuuj"
+		"editQuestionWithCheckpoint": "Edycja tej wiadomości spowoduje usunięcie wszystkich późniejszych wiadomości w rozmowie. Czy chcesz również cofnąć wszystkie zmiany do tego punktu kontrolnego?",
+		"deleteQuestionWithCheckpoint": "Usunięcie tej wiadomości spowoduje usunięcie wszystkich późniejszych wiadomości w rozmowie. Czy chcesz również cofnąć wszystkie zmiany do tego punktu kontrolnego?",
+		"editOnly": "Nie, tylko edytuj wiadomość",
+		"deleteOnly": "Nie, tylko usuń wiadomość",
+		"restoreToCheckpoint": "Tak, przywróć punkt kontrolny",
+		"proceed": "Kontynuuj",
+		"dontShowAgain": "Nie pokazuj ponownie"
 	},
 	"time_ago": {
 		"just_now": "przed chwilą",

webview-ui/src/i18n/locales/pt-BR/common.json

@@ -72,7 +72,13 @@
 		"deleteWarning": "Excluir esta mensagem irá excluir todas as mensagens subsequentes na conversa. Deseja prosseguir?",
 		"editMessage": "Editar Mensagem",
 		"editWarning": "Editar esta mensagem irá excluir todas as mensagens subsequentes na conversa. Deseja prosseguir?",
-		"proceed": "Prosseguir"
+		"editQuestionWithCheckpoint": "Editar esta mensagem irá excluir todas as mensagens posteriores na conversa. Você também deseja desfazer todas as alterações até este checkpoint?",
+		"deleteQuestionWithCheckpoint": "Excluir esta mensagem irá excluir todas as mensagens posteriores na conversa. Você também deseja desfazer todas as alterações até este checkpoint?",
+		"editOnly": "Não, apenas editar a mensagem",
+		"deleteOnly": "Não, apenas excluir a mensagem",
+		"restoreToCheckpoint": "Sim, restaurar o checkpoint",
+		"proceed": "Prosseguir",
+		"dontShowAgain": "Não mostrar novamente"
 	},
 	"time_ago": {
 		"just_now": "agora mesmo",

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

@@ -72,7 +72,13 @@
 		"deleteWarning": "Удаление этого сообщения приведет к удалению всех последующих сообщений в разговоре. Хотите продолжить?",
 		"editMessage": "Редактировать Сообщение",
 		"editWarning": "Редактирование этого сообщения приведет к удалению всех последующих сообщений в разговоре. Хотите продолжить?",
-		"proceed": "Продолжить"
+		"editQuestionWithCheckpoint": "Редактирование этого сообщения приведет к удалению всех последующих сообщений в разговоре. Хотите также отменить все изменения до этой контрольной точки?",
+		"deleteQuestionWithCheckpoint": "Удаление этого сообщения приведет к удалению всех последующих сообщений в разговоре. Хотите также отменить все изменения до этой контрольной точки?",
+		"editOnly": "Нет, только редактировать сообщение",
+		"deleteOnly": "Нет, только удалить сообщение",
+		"restoreToCheckpoint": "Да, восстановить контрольную точку",
+		"proceed": "Продолжить",
+		"dontShowAgain": "Больше не показывать"
 	},
 	"time_ago": {
 		"just_now": "только что",

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

@@ -72,7 +72,13 @@
 		"deleteWarning": "Bu mesajı silmek, konuşmadaki sonraki tüm mesajları da silecektir. Devam etmek istiyor musun?",
 		"editMessage": "Mesajı Düzenle",
 		"editWarning": "Bu mesajı düzenlemek, konuşmadaki sonraki tüm mesajları da silecektir. Devam etmek istiyor musun?",
-		"proceed": "Devam Et"
+		"editQuestionWithCheckpoint": "Bu mesajı düzenlemek, konuşmadaki sonraki tüm mesajları da silecektir. Bu kontrol noktasına kadar olan tüm değişiklikleri de geri almak istiyor musun?",
+		"deleteQuestionWithCheckpoint": "Bu mesajı silmek, konuşmadaki sonraki tüm mesajları da silecektir. Bu kontrol noktasına kadar olan tüm değişiklikleri de geri almak istiyor musun?",
+		"editOnly": "Hayır, sadece mesajı düzenle",
+		"deleteOnly": "Hayır, sadece mesajı sil",
+		"restoreToCheckpoint": "Evet, kontrol noktasını geri yükle",
+		"proceed": "Devam Et",
+		"dontShowAgain": "Tekrar gösterme"
 	},
 	"time_ago": {
 		"just_now": "şimdi",

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

@@ -72,7 +72,13 @@
 		"deleteWarning": "Xóa tin nhắn này sẽ xóa tất cả các tin nhắn tiếp theo trong cuộc trò chuyện. Bạn có muốn tiếp tục không?",
 		"editMessage": "Chỉnh Sửa Tin Nhắn",
 		"editWarning": "Chỉnh sửa tin nhắn này sẽ xóa tất cả các tin nhắn tiếp theo trong cuộc trò chuyện. Bạn có muốn tiếp tục không?",
-		"proceed": "Tiếp Tục"
+		"editQuestionWithCheckpoint": "Chỉnh sửa tin nhắn này sẽ xóa tất cả các tin nhắn sau đó trong cuộc trò chuyện. Bạn có muốn hoàn tác tất cả các thay đổi về checkpoint này không?",
+		"deleteQuestionWithCheckpoint": "Xóa tin nhắn này sẽ xóa tất cả các tin nhắn sau đó trong cuộc trò chuyện. Bạn có muốn hoàn tác tất cả các thay đổi về checkpoint này không?",
+		"editOnly": "Không, chỉ chỉnh sửa tin nhắn",
+		"deleteOnly": "Không, chỉ xóa tin nhắn",
+		"restoreToCheckpoint": "Có, khôi phục checkpoint",
+		"proceed": "Tiếp Tục",
+		"dontShowAgain": "Không hiển thị lại"
 	},
 	"time_ago": {
 		"just_now": "vừa xong",

webview-ui/src/i18n/locales/zh-CN/common.json

@@ -72,7 +72,13 @@
 		"deleteWarning": "删除此消息将删除对话中的所有后续消息。是否继续？",
 		"editMessage": "编辑消息",
 		"editWarning": "编辑此消息将删除对话中的所有后续消息。是否继续？",
-		"proceed": "继续"
+		"editQuestionWithCheckpoint": "编辑此消息将删除对话中的所有后续消息。是否同时将所有变更撤销到此存档点？",
+		"deleteQuestionWithCheckpoint": "删除此消息将删除对话中的所有后续消息。是否同时将所有变更撤销到此存档点？",
+		"editOnly": "否，仅编辑消息",
+		"deleteOnly": "否，仅删除消息",
+		"restoreToCheckpoint": "是，恢复存档点",
+		"proceed": "继续",
+		"dontShowAgain": "不再显示"
 	},
 	"time_ago": {
 		"just_now": "刚刚",

webview-ui/src/i18n/locales/zh-TW/common.json

@@ -71,8 +71,14 @@
 		"deleteMessage": "刪除訊息",
 		"deleteWarning": "刪除此訊息將會刪除對話中所有後續的訊息。您要繼續嗎？",
 		"editMessage": "編輯訊息",
-		"editWarning": "編輯此訊息將會刪除對話中所有後續的訊息。您要繼續嗎？",
-		"proceed": "繼續"
+		"editWarning": "編輯此訊息將刪除對話中的所有後續訊息。是否繼續？",
+		"editQuestionWithCheckpoint": "編輯此訊息將刪除對話中的所有後續訊息。是否同時將所有變更撤銷到此存檔點？",
+		"deleteQuestionWithCheckpoint": "刪除此訊息將刪除對話中的所有後續訊息。是否同時將所有變更撤銷到此存檔點？",
+		"editOnly": "否，僅編輯訊息",
+		"deleteOnly": "否，僅刪除訊息",
+		"restoreToCheckpoint": "是，恢復存檔點",
+		"proceed": "繼續",
+		"dontShowAgain": "不再顯示"
 	},
 	"time_ago": {
 		"just_now": "剛剛",