fix: initial checkpoint restore now uses checkpointRestoreToBase

The initial checkpoint has no message in clineMessages (ts=0 doesn't exist),
so checkpointRestore was returning early without doing anything.

Now when isInitial=true is passed in the restore payload, the handler
calls checkpointRestoreToBase() instead which properly restores to the
baseHash without needing to find a message.

Author: hannesrudolph

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

@@ -737,6 +737,7 @@ export const checkoutRestorePayloadSchema = z.object({
 	ts: z.number(),
 	commitHash: z.string(),
 	mode: z.enum(["preview", "restore"]),
+	isInitial: z.boolean().optional(),
 })
 
 export type CheckpointRestorePayload = z.infer<typeof checkoutRestorePayloadSchema>

src/core/webview/webviewMessageHandler.ts

@@ -1200,7 +1200,18 @@ export const webviewMessageHandler = async (
 				}
 
 				try {
-					await provider.getCurrentTask()?.checkpointRestore(result.data)
+					// For initial checkpoint, use checkpointRestoreToBase instead of checkpointRestore
+					// because there's no message with ts to find in clineMessages
+					if (result.data.isInitial) {
+						const { checkpointRestoreToBase } = await import("../checkpoints")
+						const success = await checkpointRestoreToBase(provider.getCurrentTask()!)
+
+						if (!success) {
+							vscode.window.showErrorMessage(t("common:errors.checkpoint_restore_base_failed"))
+						}
+					} else {
+						await provider.getCurrentTask()?.checkpointRestore(result.data)
+					}
 				} catch (error) {
 					vscode.window.showErrorMessage(t("common:errors.checkpoint_failed"))
 				}

webview-ui/src/components/chat/checkpoints/CheckpointMenu.tsx

@@ -75,14 +75,20 @@ export const CheckpointMenu = ({ ts, commitHash, checkpoint, isInitial, onOpenCh
 	}, [ts, commitHash])
 
 	const onPreview = useCallback(() => {
-		vscode.postMessage({ type: "checkpointRestore", payload: { ts, commitHash, mode: "preview" } })
+		vscode.postMessage({
+			type: "checkpointRestore",
+			payload: { ts, commitHash, mode: "preview", ...(isInitial && { isInitial: true }) },
+		})
 		setRestoreOpen(false)
-	}, [ts, commitHash, setRestoreOpen])
+	}, [ts, commitHash, isInitial, setRestoreOpen])
 
 	const onRestore = useCallback(() => {
-		vscode.postMessage({ type: "checkpointRestore", payload: { ts, commitHash, mode: "restore" } })
+		vscode.postMessage({
+			type: "checkpointRestore",
+			payload: { ts, commitHash, mode: "restore", ...(isInitial && { isInitial: true }) },
+		})
 		setRestoreOpen(false)
-	}, [ts, commitHash, setRestoreOpen])
+	}, [ts, commitHash, isInitial, setRestoreOpen])
 
 	const handleOpenChange = useCallback(
 		(open: boolean) => {