From a89acc58ed47ec80768d72b3d53691fd3650b9ee Mon Sep 17 00:00:00 2001
From: Hannes Rudolph <hrudolph@gmail.com>
Date: Thu, 29 Jan 2026 13:49:31 -0700
Subject: [PATCH] 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.
---
 packages/types/src/vscode-extension-host.ts        |  1 +
 src/core/webview/webviewMessageHandler.ts          | 13 ++++++++++++-
 .../components/chat/checkpoints/CheckpointMenu.tsx | 14 ++++++++++----
 3 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/packages/types/src/vscode-extension-host.ts b/packages/types/src/vscode-extension-host.ts
index 77c8597a28..f88cf9fcc4 100644
--- a/packages/types/src/vscode-extension-host.ts
+++ b/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>
diff --git a/src/core/webview/webviewMessageHandler.ts b/src/core/webview/webviewMessageHandler.ts
index 8246dda472..de272773c9 100644
--- a/src/core/webview/webviewMessageHandler.ts
+++ b/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"))
 				}
diff --git a/webview-ui/src/components/chat/checkpoints/CheckpointMenu.tsx b/webview-ui/src/components/chat/checkpoints/CheckpointMenu.tsx
index d3f0a7f343..d3a7730114 100644
--- a/webview-ui/src/components/chat/checkpoints/CheckpointMenu.tsx
+++ b/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) => {