feat: add restore to task start button

packages/types/src/vscode-extension-host.ts

@@ -524,6 +524,7 @@ export interface WebviewMessage {
 		| "openCustomModesSettings"
 		| "checkpointDiff"
 		| "checkpointRestore"
+		| "restoreToTaskStart"
 		| "deleteMcpServer"
 		| "codebaseIndexEnabled"
 		| "telemetrySetting"

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

@@ -1,7 +1,13 @@
 import { describe, it, expect, vi, beforeEach, afterEach, Mock } from "vitest"
 import { Task } from "../../task/Task"
 import { ClineProvider } from "../../webview/ClineProvider"
-import { checkpointSave, checkpointRestore, checkpointDiff, getCheckpointService } from "../index"
+import {
+	checkpointSave,
+	checkpointRestore,
+	checkpointRestoreToBase,
+	checkpointDiff,
+	getCheckpointService,
+} from "../index"
 import { MessageManager } from "../../message-manager"
 import * as vscode from "vscode"
 
@@ -296,6 +302,56 @@ describe("Checkpoint functionality", () => {
 		})
 	})
 
+	describe("checkpointRestoreToBase", () => {
+		beforeEach(() => {
+			mockCheckpointService.baseHash = "initial-commit-hash"
+		})
+
+		it("should restore to base hash successfully", async () => {
+			const result = await checkpointRestoreToBase(mockTask)
+
+			expect(result).toBe(true)
+			expect(mockCheckpointService.restoreCheckpoint).toHaveBeenCalledWith("initial-commit-hash")
+			expect(mockProvider.postMessageToWebview).toHaveBeenCalledWith({
+				type: "currentCheckpointUpdated",
+				text: "initial-commit-hash",
+			})
+			expect(mockProvider.cancelTask).toHaveBeenCalled()
+		})
+
+		it("should return false if no checkpoint service available", async () => {
+			mockTask.checkpointService = undefined
+			mockTask.enableCheckpoints = false
+
+			const result = await checkpointRestoreToBase(mockTask)
+
+			expect(result).toBe(false)
+			expect(mockCheckpointService.restoreCheckpoint).not.toHaveBeenCalled()
+		})
+
+		it("should return false if no baseHash available", async () => {
+			mockCheckpointService.baseHash = undefined
+
+			const result = await checkpointRestoreToBase(mockTask)
+
+			expect(result).toBe(false)
+			expect(mockCheckpointService.restoreCheckpoint).not.toHaveBeenCalled()
+			expect(mockProvider.log).toHaveBeenCalledWith("[checkpointRestoreToBase] no baseHash available")
+		})
+
+		it("should disable checkpoints on error", async () => {
+			mockCheckpointService.restoreCheckpoint.mockRejectedValue(new Error("Restore failed"))
+
+			const result = await checkpointRestoreToBase(mockTask)
+
+			expect(result).toBe(false)
+			expect(mockTask.enableCheckpoints).toBe(false)
+			expect(mockProvider.log).toHaveBeenCalledWith(
+				"[checkpointRestoreToBase] disabling checkpoints for this task",
+			)
+		})
+	})
+
 	describe("checkpointDiff", () => {
 		beforeEach(() => {
 			mockTask.clineMessages = [

src/core/checkpoints/index.ts

@@ -301,6 +301,46 @@ export async function checkpointRestore(
 	}
 }
 
+/**
+ * Restore the workspace to its initial state (baseHash) - the state when the shadow git repo was initialized.
+ * This is a simpler version of checkpointRestore that doesn't need to rewind messages since we're
+ * restoring to the very beginning of the task.
+ * @returns true if restoration was successful, false otherwise
+ */
+export async function checkpointRestoreToBase(task: Task): Promise<boolean> {
+	const service = await getCheckpointService(task)
+
+	if (!service) {
+		return false
+	}
+
+	const baseHash = service.baseHash
+
+	if (!baseHash) {
+		const provider = task.providerRef.deref()
+		provider?.log("[checkpointRestoreToBase] no baseHash available")
+		return false
+	}
+
+	const provider = task.providerRef.deref()
+
+	try {
+		await service.restoreCheckpoint(baseHash)
+		TelemetryService.instance.captureCheckpointRestored(task.taskId)
+		await provider?.postMessageToWebview({ type: "currentCheckpointUpdated", text: baseHash })
+
+		// Cancel the task to reinitialize with the restored state
+		// This follows the same pattern as checkpointRestore
+		provider?.cancelTask()
+
+		return true
+	} catch (err) {
+		provider?.log("[checkpointRestoreToBase] disabling checkpoints for this task")
+		task.enableCheckpoints = false
+		return false
+	}
+}
+
 export type CheckpointDiffOptions = {
 	ts?: number
 	previousCommitHash?: string

src/core/webview/webviewMessageHandler.ts

@@ -1212,6 +1212,42 @@ export const webviewMessageHandler = async (
 
 			break
 		}
+		case "restoreToTaskStart": {
+			const currentTask = provider.getCurrentTask()
+
+			if (!currentTask) {
+				vscode.window.showErrorMessage(t("common:errors.checkpoint_no_active_task"))
+				break
+			}
+
+			if (!currentTask.enableCheckpoints) {
+				vscode.window.showErrorMessage(t("common:errors.checkpoint_not_enabled"))
+				break
+			}
+
+			// Cancel the current task first
+			await provider.cancelTask()
+
+			try {
+				await pWaitFor(() => provider.getCurrentTask()?.isInitialized === true, { timeout: 3_000 })
+			} catch (error) {
+				vscode.window.showErrorMessage(t("common:errors.checkpoint_timeout"))
+				break
+			}
+
+			try {
+				const { checkpointRestoreToBase } = await import("../checkpoints")
+				const success = await checkpointRestoreToBase(provider.getCurrentTask()!)
+
+				if (!success) {
+					vscode.window.showErrorMessage(t("common:errors.checkpoint_restore_base_failed"))
+				}
+			} catch (error) {
+				vscode.window.showErrorMessage(t("common:errors.checkpoint_failed"))
+			}
+
+			break
+		}
 		case "cancelTask":
 			await provider.cancelTask()
 			break

src/i18n/locales/en/common.json

@@ -28,6 +28,9 @@
 		"could_not_open_file_generic": "Could not open file!",
 		"checkpoint_timeout": "Timed out when attempting to restore checkpoint.",
 		"checkpoint_failed": "Failed to restore checkpoint.",
+		"checkpoint_no_active_task": "No active task to restore.",
+		"checkpoint_not_enabled": "Checkpoints are not enabled for this task.",
+		"checkpoint_restore_base_failed": "Failed to restore workspace to initial state.",
 		"git_not_installed": "Git is required for the checkpoints feature. Please install Git to enable checkpoints.",
 		"checkpoint_no_first": "No first checkpoint to compare.",
 		"checkpoint_no_previous": "No previous checkpoint to compare.",