diff --git a/src/core/webview/__tests__/checkpointRestoreHandler.test.ts b/src/core/webview/__tests__/checkpointRestoreHandler.test.ts new file mode 100644 index 0000000000..d79349ed63 --- /dev/null +++ b/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", + ) + }) + }) + }) +}) diff --git a/src/core/webview/checkpointRestoreHandler.ts b/src/core/webview/checkpointRestoreHandler.ts new file mode 100644 index 0000000000..a9a6c83f66 --- /dev/null +++ b/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 { + 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 { + try { + await pWaitFor(() => provider.getCurrentCline()?.isInitialized === true, { + timeout: timeoutMs, + }) + return true + } catch (error) { + vscode.window.showErrorMessage(t("common:errors.checkpoint_timeout")) + return false + } +} diff --git a/src/core/webview/webviewMessageHandler.ts b/src/core/webview/webviewMessageHandler.ts index 682c46cc5e..0258fe1657 100644 --- a/src/core/webview/webviewMessageHandler.ts +++ b/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