refactor

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

@@ -0,0 +1,164 @@
+import { describe, it, expect, vi, beforeEach } from "vitest"
+import { handleCheckpointRestoreOperation, hasValidCheckpoint } from "../checkpointRestoreHandler"
+import { saveTaskMessages } from "../../task-persistence"
+import * as vscode from "vscode"
+
+vi.mock("../../task-persistence", () => ({
+	saveTaskMessages: vi.fn(),
+}))
+
+vi.mock("vscode", () => ({
+	window: {
+		showErrorMessage: vi.fn(),
+	},
+}))
+
+describe("checkpointRestoreHandler", () => {
+	let mockProvider: any
+	let mockCline: any
+
+	beforeEach(() => {
+		vi.clearAllMocks()
+
+		mockProvider = {
+			contextProxy: {
+				globalStorageUri: {
+					fsPath: "/test/global/storage",
+				},
+			},
+			getTaskWithId: vi.fn().mockResolvedValue({
+				historyItem: { id: "test-task", messages: [] },
+			}),
+			initClineWithHistoryItem: vi.fn(),
+		}
+
+		mockCline = {
+			taskId: "test-task",
+			clineMessages: [
+				{ ts: 1, text: "Message 1" },
+				{ ts: 2, text: "Message 2", checkpoint: { hash: "abc123" } },
+				{ ts: 3, text: "Message 3" },
+			],
+			checkpointRestore: vi.fn(),
+		}
+	})
+
+	describe("hasValidCheckpoint", () => {
+		it("should return true for valid checkpoint", () => {
+			const message = { checkpoint: { hash: "abc123" } }
+			expect(hasValidCheckpoint(message)).toBe(true)
+		})
+
+		it("should return false for missing checkpoint", () => {
+			const message = { text: "No checkpoint" }
+			expect(hasValidCheckpoint(message)).toBe(false)
+		})
+
+		it("should return false for invalid checkpoint structure", () => {
+			expect(hasValidCheckpoint({ checkpoint: "invalid" })).toBe(false)
+			expect(hasValidCheckpoint({ checkpoint: {} })).toBe(false)
+			expect(hasValidCheckpoint({ checkpoint: { hash: 123 } })).toBe(false)
+		})
+	})
+
+	describe("handleCheckpointRestoreOperation", () => {
+		describe("delete operation", () => {
+			it("should handle delete operation correctly", async () => {
+				await handleCheckpointRestoreOperation({
+					provider: mockProvider,
+					currentCline: mockCline,
+					messageTs: 2,
+					messageIndex: 1,
+					checkpoint: { hash: "abc123" },
+					operation: "delete",
+				})
+
+				// Should call checkpointRestore with correct params
+				expect(mockCline.checkpointRestore).toHaveBeenCalledWith({
+					ts: 2,
+					commitHash: "abc123",
+					mode: "restore",
+					operation: "delete",
+				})
+
+				// Should save messages after restoration
+				expect(saveTaskMessages).toHaveBeenCalledWith({
+					messages: mockCline.clineMessages,
+					taskId: "test-task",
+					globalStoragePath: "/test/global/storage",
+				})
+
+				// Should reinitialize the task
+				expect(mockProvider.getTaskWithId).toHaveBeenCalledWith("test-task")
+				expect(mockProvider.initClineWithHistoryItem).toHaveBeenCalledWith({
+					id: "test-task",
+					messages: [],
+				})
+			})
+		})
+
+		describe("edit operation", () => {
+			it("should handle edit operation correctly", async () => {
+				const editData = {
+					editedContent: "Edited content",
+					images: ["image1.png"],
+					apiConversationHistoryIndex: 1,
+				}
+
+				await handleCheckpointRestoreOperation({
+					provider: mockProvider,
+					currentCline: mockCline,
+					messageTs: 2,
+					messageIndex: 1,
+					checkpoint: { hash: "abc123" },
+					operation: "edit",
+					editData,
+				})
+
+				// Should set pendingEditAfterRestore on provider
+				expect(mockProvider.pendingEditAfterRestore).toEqual({
+					messageTs: 2,
+					editedContent: "Edited content",
+					images: ["image1.png"],
+					messageIndex: 1,
+					apiConversationHistoryIndex: 1,
+					originalCheckpoint: { hash: "abc123" },
+				})
+
+				// Should call checkpointRestore with correct params
+				expect(mockCline.checkpointRestore).toHaveBeenCalledWith({
+					ts: 2,
+					commitHash: "abc123",
+					mode: "restore",
+					operation: "edit",
+				})
+
+				// Should NOT save messages or reinitialize for edit
+				expect(saveTaskMessages).not.toHaveBeenCalled()
+				expect(mockProvider.initClineWithHistoryItem).not.toHaveBeenCalled()
+			})
+		})
+
+		describe("error handling", () => {
+			it("should handle errors and show error message", async () => {
+				const error = new Error("Checkpoint restore failed")
+				mockCline.checkpointRestore.mockRejectedValue(error)
+
+				await expect(
+					handleCheckpointRestoreOperation({
+						provider: mockProvider,
+						currentCline: mockCline,
+						messageTs: 2,
+						messageIndex: 1,
+						checkpoint: { hash: "abc123" },
+						operation: "delete",
+					}),
+				).rejects.toThrow("Checkpoint restore failed")
+
+				expect(vscode.window.showErrorMessage).toHaveBeenCalledWith(
+					"Error during checkpoint restore: Checkpoint restore failed",
+				)
+			})
+		})
+	})
+})

src/core/webview/checkpointRestoreHandler.ts

@@ -0,0 +1,103 @@
+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 edit operations, set up pending edit data before restoration
+		if (operation === "edit" && editData) {
+			;(provider as any).pendingEditAfterRestore = {
+				messageTs,
+				editedContent: editData.editedContent,
+				images: editData.images,
+				messageIndex: config.messageIndex,
+				apiConversationHistoryIndex: editData.apiConversationHistoryIndex,
+				originalCheckpoint: checkpoint,
+			}
+		}
+
+		// 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.initClineWithHistoryItem(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
+	}
+}
+
+/**
+ * Validates if a message has a valid checkpoint for restoration
+ */
+export function hasValidCheckpoint(message: any): boolean {
+	return (
+		(message?.checkpoint &&
+			typeof message.checkpoint === "object" &&
+			"hash" in message.checkpoint &&
+			typeof message.checkpoint.hash === "string") ||
+		false
+	)
+}
+
+/**
+ * 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.getCurrentCline()?.isInitialized === true, {
+			timeout: timeoutMs,
+		})
+		return true
+	} catch (error) {
+		vscode.window.showErrorMessage(t("common:errors.checkpoint_timeout"))
+		return false
+	}
+}

src/core/webview/webviewMessageHandler.ts

@@ -19,6 +19,11 @@ import { type ApiMessage } from "../task-persistence/apiMessages"
 import { saveTaskMessages } from "../task-persistence"
 
 import { ClineProvider } from "./ClineProvider"
+import {
+	handleCheckpointRestoreOperation,
+	hasValidCheckpoint,
+	waitForClineInitialization,
+} from "./checkpointRestoreHandler"
 import { changeLanguage, t } from "../../i18n"
 import { Package } from "../../shared/package"
 import { RouterName, toRouterName, ModelRecord } from "../../shared/api"
@@ -154,29 +159,18 @@ export const webviewMessageHandler = async (
 
 			if (messageIndex !== -1) {
 				try {
-					// If checkpoint restoration is requested, restore to the checkpoint first
-					if (restoreCheckpoint) {
-						const targetMessage = currentCline.clineMessages[messageIndex]
-						if (
-							targetMessage?.checkpoint &&
-							typeof targetMessage.checkpoint === "object" &&
-							"hash" in targetMessage.checkpoint
-						) {
-							await currentCline.checkpointRestore({
-								ts: targetMessage.ts!,
-								commitHash: targetMessage.checkpoint.hash as string,
-								mode: "restore",
-								operation: "delete",
-							})
+					const targetMessage = currentCline.clineMessages[messageIndex]
 
-							// Save the updated messages to disk after checkpoint restoration
-							// This ensures the deleted messages are persisted before reinitialization
-							await saveTaskMessages({
-								messages: currentCline.clineMessages,
-								taskId: currentCline.taskId,
-								globalStoragePath: provider.contextProxy.globalStorageUri.fsPath,
-							})
-						}
+					// If checkpoint restoration is requested, restore to the checkpoint first
+					if (restoreCheckpoint && hasValidCheckpoint(targetMessage)) {
+						await handleCheckpointRestoreOperation({
+							provider,
+							currentCline,
+							messageTs: targetMessage.ts!,
+							messageIndex,
+							checkpoint: targetMessage.checkpoint as { hash: string },
+							operation: "delete",
+						})
 					} else {
 						// For non-checkpoint deletes, preserve checkpoint associations for remaining messages
 						// Store checkpoints from messages that will be preserved
@@ -206,13 +200,6 @@ export const webviewMessageHandler = async (
 							globalStoragePath: provider.contextProxy.globalStorageUri.fsPath,
 						})
 					}
-
-					const { historyItem } = await provider.getTaskWithId(currentCline.taskId)
-
-					// Initialize with history item after deletion (only for checkpoint restores)
-					if (restoreCheckpoint) {
-						await provider.initClineWithHistoryItem(historyItem)
-					}
 				} catch (error) {
 					console.error("Error in delete message:", error)
 					vscode.window.showErrorMessage(
@@ -307,36 +294,23 @@ export const webviewMessageHandler = async (
 					const originalCheckpoint = targetMessage?.checkpoint
 
 					// If checkpoint restoration is requested, restore to the checkpoint first
-					if (restoreCheckpoint) {
-						if (
-							originalCheckpoint &&
-							typeof originalCheckpoint === "object" &&
-							"hash" in originalCheckpoint
-						) {
-							// Store the edited content and images for after restoration
-							const editData = { text: editedContent, images }
-
-							// Set a flag on the provider to indicate we need to process an edit after restoration
-							;(provider as any).pendingEditAfterRestore = {
-								messageTs,
+					if (restoreCheckpoint && hasValidCheckpoint(targetMessage)) {
+						await handleCheckpointRestoreOperation({
+							provider,
+							currentCline,
+							messageTs: targetMessage.ts!,
+							messageIndex,
+							checkpoint: targetMessage.checkpoint as { hash: string },
+							operation: "edit",
+							editData: {
 								editedContent,
 								images,
-								messageIndex,
 								apiConversationHistoryIndex,
-								originalCheckpoint, // Preserve the checkpoint for the new message
-							}
-
-							await currentCline.checkpointRestore({
-								ts: targetMessage.ts!,
-								commitHash: originalCheckpoint.hash as string,
-								mode: "restore",
-								operation: "edit",
-							})
-
-							// The task will be cancelled and reinitialized by checkpointRestore
-							// The pending edit will be processed in the reinitialized task
-							return
-						}
+							},
+						})
+						// The task will be cancelled and reinitialized by checkpointRestore
+						// The pending edit will be processed in the reinitialized task
+						return
 					}
 
 					// For non-checkpoint edits, preserve checkpoint associations for remaining messages