feat: add restore to task start button

Adds a button in the task header that allows users to restore the workspace
to its initial state when the task was created. This uses the baseHash
checkpoint from the shadow git repository.

- Add checkpointRestoreToBase() function in checkpoints/index.ts
- Add restoreToTaskStart message handler in webviewMessageHandler.ts
- Add RestoreTaskDialog component with confirmation dialog
- Add restore button to TaskActions (visible when checkpoints enabled)
- Add translations for all 17 supported locales
- Add 4 unit tests for the new checkpoint restore function

Author: hannesrudolph

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/ca/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.",