From ab8376fee495c7bf33392af197a245bd66ad4bce Mon Sep 17 00:00:00 2001
From: Saoud Rizwan <7799382+saoudrizwan@users.noreply.github.com>
Date: Wed, 19 Feb 2025 18:57:16 -0800
Subject: [PATCH] Add 'Restore Task' option back to checkpoints

---
 src/core/Cline.ts                              | 18 ++++++++++--------
 .../src/components/common/CheckmarkControl.tsx |  4 ++--
 2 files changed, 12 insertions(+), 10 deletions(-)

diff --git a/src/core/Cline.ts b/src/core/Cline.ts
index 9209eeacb5..b1c6baa90c 100644
--- a/src/core/Cline.ts
+++ b/src/core/Cline.ts
@@ -353,15 +353,17 @@ export class Cline {
 					break
 			}
 
-			// Set isCheckpointCheckedOut flag on the message
-			// Find all checkpoint messages before this one
-			const checkpointMessages = this.clineMessages.filter((m) => m.say === "checkpoint_created")
-			const currentMessageIndex = checkpointMessages.findIndex((m) => m.ts === messageTs)
+			if (restoreType !== "task") {
+				// Set isCheckpointCheckedOut flag on the message
+				// Find all checkpoint messages before this one
+				const checkpointMessages = this.clineMessages.filter((m) => m.say === "checkpoint_created")
+				const currentMessageIndex = checkpointMessages.findIndex((m) => m.ts === messageTs)
 
-			// Set isCheckpointCheckedOut to false for all checkpoint messages
-			checkpointMessages.forEach((m, i) => {
-				m.isCheckpointCheckedOut = i === currentMessageIndex
-			})
+				// Set isCheckpointCheckedOut to false for all checkpoint messages
+				checkpointMessages.forEach((m, i) => {
+					m.isCheckpointCheckedOut = i === currentMessageIndex
+				})
+			}
 
 			await this.saveClineMessages()
 
diff --git a/webview-ui/src/components/common/CheckmarkControl.tsx b/webview-ui/src/components/common/CheckmarkControl.tsx
index b57a6a0bf6..9641e77b7d 100644
--- a/webview-ui/src/components/common/CheckmarkControl.tsx
+++ b/webview-ui/src/components/common/CheckmarkControl.tsx
@@ -180,7 +180,7 @@ export const CheckmarkControl = ({ messageTs, isCheckpointCheckedOut }: Checkmar
 										what will be reverted)
 									</p>
 								</RestoreOption>
-								{/* <RestoreOption>
+								<RestoreOption>
 									<VSCodeButton
 										onClick={handleRestoreTask}
 										disabled={restoreTaskDisabled}
@@ -192,7 +192,7 @@ export const CheckmarkControl = ({ messageTs, isCheckpointCheckedOut }: Checkmar
 										Restore Task Only
 									</VSCodeButton>
 									<p>Deletes messages after this point (does not affect workspace files)</p>
-								</RestoreOption> */}
+								</RestoreOption>
 								<RestoreOption>
 									<VSCodeButton
 										onClick={handleRestoreBoth}