fix: checkpoint restore uses current task service before cancellation

Fixes #10402

The checkpoint restore was failing because it called cancelTask() BEFORE
calling checkpointRestore(). This destroyed the current task with its
initialized checkpoint service and created a new task. The new task had
no checkpoint service ready yet, so getCheckpointService() returned
undefined and the restore silently failed.

The fix is to call checkpointRestore() directly on the current task,
which already has its checkpoint service initialized. The checkpointRestore
function already handles calling cancelTask() internally at the end to
reinitialize the UI.

Changes:
- Remove premature cancelTask() call in checkpointRestore message handler
- Call checkpointRestore() on current task directly
- Remove unused pWaitFor import

src/core/webview/webviewMessageHandler.ts

@@ -3,7 +3,6 @@ import * as path from "path"
 import * as os from "os"
 import * as fs from "fs/promises"
 import { getRooDirectoriesForCwd } from "../../services/roo-config/index.js"
-import pWaitFor from "p-wait-for"
 import * as vscode from "vscode"
 
 import {
@@ -1101,23 +1100,29 @@ export const webviewMessageHandler = async (
 			break
 		case "checkpointRestore": {
 			const result = checkoutRestorePayloadSchema.safeParse(message.payload)
-
+	
 			if (result.success) {
-				await provider.cancelTask()
-
-				try {
-					await pWaitFor(() => provider.getCurrentTask()?.isInitialized === true, { timeout: 3_000 })
-				} catch (error) {
+				// Get the current task BEFORE any cancellation.
+				// The checkpoint service is only available on the current task.
+				const currentTask = provider.getCurrentTask()
+	
+				if (!currentTask) {
 					vscode.window.showErrorMessage(t("common:errors.checkpoint_timeout"))
+					break
 				}
-
+	
 				try {
-					await provider.getCurrentTask()?.checkpointRestore(result.data)
+					// Call checkpointRestore on the CURRENT task.
+					// This will:
+					// 1. Use the current task's checkpoint service to restore files
+					// 2. Truncate messages appropriately
+					// 3. Call cancelTask() internally to reinitialize the UI
+					await currentTask.checkpointRestore(result.data)
 				} catch (error) {
 					vscode.window.showErrorMessage(t("common:errors.checkpoint_failed"))
 				}
 			}
-
+	
 			break
 		}
 		case "cancelTask":