From fce7f1bea3ad95a12e0e2b8a45c53a2d938cf87c Mon Sep 17 00:00:00 2001
From: Saoud Rizwan <7799382+saoudrizwan@users.noreply.github.com>
Date: Sun, 9 Feb 2025 03:03:09 -0800
Subject: [PATCH] Use more visual checkpoints tracking

---
 package-lock.json                             |   4 +-
 src/core/Cline.ts                             | 215 ++++++----
 src/shared/ExtensionMessage.ts                |   2 +
 webview-ui/package-lock.json                  |  60 +++
 webview-ui/package.json                       |   3 +-
 .../src/components/chat/BrowserSessionRow.tsx |  31 +-
 webview-ui/src/components/chat/ChatRow.tsx    |  11 +-
 webview-ui/src/components/chat/TaskHeader.tsx |   2 +-
 .../components/common/CheckmarkControl.tsx    | 394 ++++++++++++++++++
 9 files changed, 609 insertions(+), 113 deletions(-)
 create mode 100644 webview-ui/src/components/common/CheckmarkControl.tsx

diff --git a/package-lock.json b/package-lock.json
index d193735592..3591d6bd3b 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -1,12 +1,12 @@
 {
 	"name": "claude-dev",
-	"version": "3.2.13",
+	"version": "3.3.1",
 	"lockfileVersion": 3,
 	"requires": true,
 	"packages": {
 		"": {
 			"name": "claude-dev",
-			"version": "3.2.13",
+			"version": "3.3.1",
 			"license": "Apache-2.0",
 			"dependencies": {
 				"@anthropic-ai/bedrock-sdk": "^0.10.2",
diff --git a/src/core/Cline.ts b/src/core/Cline.ts
index 9f5af06f3f..81f5b13d11 100644
--- a/src/core/Cline.ts
+++ b/src/core/Cline.ts
@@ -353,6 +353,18 @@ 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)
+
+			// Set isCheckpointCheckedOut to false for all checkpoint messages
+			checkpointMessages.forEach((m, i) => {
+				m.isCheckpointCheckedOut = i === currentMessageIndex
+			})
+
+			await this.saveClineMessages()
+
 			await this.providerRef.deref()?.postMessageToWebview({ type: "relinquishControl" })
 
 			this.providerRef.deref()?.cancelTask() // the task is already cancelled by the provider beforehand, but we need to re-init to get the updated messages
@@ -1062,40 +1074,67 @@ export class Cline {
 
 	// Checkpoints
 
-	async saveCheckpoint() {
+	async saveCheckpoint(isAttemptCompletionMessage: boolean = false) {
 		const commitHash = await this.checkpointTracker?.commit() // silently fails for now
+		// Set isCheckpointCheckedOut to false for all checkpoint_created messages
+		this.clineMessages.forEach((message) => {
+			if (message.say === "checkpoint_created") {
+				message.isCheckpointCheckedOut = false
+			}
+		})
 		if (commitHash) {
-			// Start from the end and work backwards until we find a tool use or another message with a hash
-			for (let i = this.clineMessages.length - 1; i >= 0; i--) {
-				const message = this.clineMessages[i]
-				if (message.lastCheckpointHash) {
-					// Found a message with a hash, so we can stop
-					break
+			if (!isAttemptCompletionMessage) {
+				// For non-attempt completion we just say checkpoints
+				await this.say("checkpoint_created", commitHash)
+				const lastCheckpointMessage = findLast(this.clineMessages, (m) => m.say === "checkpoint_created")
+				if (lastCheckpointMessage) {
+					lastCheckpointMessage.lastCheckpointHash = commitHash
+					await this.saveClineMessages()
 				}
-				// Update this message with a hash
-				message.lastCheckpointHash = commitHash
-
-				// We only care about adding the hash to the last tool use (we don't want to add this hash to every prior message ie for tasks pre-checkpoint)
-				const isToolUse =
-					message.say === "tool" ||
-					message.ask === "tool" ||
-					message.say === "command" ||
-					message.ask === "command" ||
-					message.say === "completion_result" ||
-					message.ask === "completion_result" ||
-					message.ask === "followup" ||
-					message.say === "use_mcp_server" ||
-					message.ask === "use_mcp_server" ||
-					message.say === "browser_action" ||
-					message.say === "browser_action_launch" ||
-					message.ask === "browser_action_launch"
-
-				if (isToolUse) {
-					break
+			} else {
+				// For attempt_completion, find the last completion_result message and set its checkpoint hash. This will be used to present the 'see new changes' button
+				const lastCompletionResultMessage = findLast(
+					this.clineMessages,
+					(m) => m.say === "completion_result" || m.ask === "completion_result",
+				)
+				if (lastCompletionResultMessage) {
+					lastCompletionResultMessage.lastCheckpointHash = commitHash
+					await this.saveClineMessages()
 				}
 			}
-			// Save the updated messages
-			await this.saveClineMessages()
+
+			// Previously we checkpointed every message, but this is excessive and unnecessary.
+			// // Start from the end and work backwards until we find a tool use or another message with a hash
+			// for (let i = this.clineMessages.length - 1; i >= 0; i--) {
+			// 	const message = this.clineMessages[i]
+			// 	if (message.lastCheckpointHash) {
+			// 		// Found a message with a hash, so we can stop
+			// 		break
+			// 	}
+			// 	// Update this message with a hash
+			// 	message.lastCheckpointHash = commitHash
+
+			// 	// We only care about adding the hash to the last tool use (we don't want to add this hash to every prior message ie for tasks pre-checkpoint)
+			// 	const isToolUse =
+			// 		message.say === "tool" ||
+			// 		message.ask === "tool" ||
+			// 		message.say === "command" ||
+			// 		message.ask === "command" ||
+			// 		message.say === "completion_result" ||
+			// 		message.ask === "completion_result" ||
+			// 		message.ask === "followup" ||
+			// 		message.say === "use_mcp_server" ||
+			// 		message.ask === "use_mcp_server" ||
+			// 		message.say === "browser_action" ||
+			// 		message.say === "browser_action_launch" ||
+			// 		message.ask === "browser_action_launch"
+
+			// 	if (isToolUse) {
+			// 		break
+			// 	}
+			// }
+			// // Save the updated messages
+			// await this.saveClineMessages()
 		}
 	}
 
@@ -1601,7 +1640,7 @@ export class Cline {
 						if (!accessAllowed) {
 							await this.say("clineignore_error", relPath)
 							pushToolResult(formatResponse.toolError(formatResponse.clineIgnoreError(relPath)))
-							await this.saveCheckpoint()
+
 							break
 						}
 
@@ -1705,21 +1744,21 @@ export class Cline {
 									this.consecutiveMistakeCount++
 									pushToolResult(await this.sayAndCreateMissingParamError(block.name, "path"))
 									await this.diffViewProvider.reset()
-									await this.saveCheckpoint()
+
 									break
 								}
 								if (block.name === "replace_in_file" && !diff) {
 									this.consecutiveMistakeCount++
 									pushToolResult(await this.sayAndCreateMissingParamError("replace_in_file", "diff"))
 									await this.diffViewProvider.reset()
-									await this.saveCheckpoint()
+
 									break
 								}
 								if (block.name === "write_to_file" && !content) {
 									this.consecutiveMistakeCount++
 									pushToolResult(await this.sayAndCreateMissingParamError("write_to_file", "content"))
 									await this.diffViewProvider.reset()
-									await this.saveCheckpoint()
+
 									break
 								}
 
@@ -1792,7 +1831,7 @@ export class Cline {
 
 									if (!didApprove) {
 										await this.diffViewProvider.revertChanges()
-										await this.saveCheckpoint()
+
 										break
 									}
 								}
@@ -1841,14 +1880,16 @@ export class Cline {
 								}
 
 								await this.diffViewProvider.reset()
+
 								await this.saveCheckpoint()
+
 								break
 							}
 						} catch (error) {
 							await handleError("writing file", error)
 							await this.diffViewProvider.revertChanges()
 							await this.diffViewProvider.reset()
-							await this.saveCheckpoint()
+
 							break
 						}
 					}
@@ -1876,7 +1917,7 @@ export class Cline {
 								if (!relPath) {
 									this.consecutiveMistakeCount++
 									pushToolResult(await this.sayAndCreateMissingParamError("read_file", "path"))
-									await this.saveCheckpoint()
+
 									break
 								}
 
@@ -1884,7 +1925,7 @@ export class Cline {
 								if (!accessAllowed) {
 									await this.say("clineignore_error", relPath)
 									pushToolResult(formatResponse.toolError(formatResponse.clineIgnoreError(relPath)))
-									await this.saveCheckpoint()
+
 									break
 								}
 
@@ -1905,19 +1946,18 @@ export class Cline {
 									this.removeLastPartialMessageIfExistsWithType("say", "tool")
 									const didApprove = await askApproval("tool", completeMessage)
 									if (!didApprove) {
-										await this.saveCheckpoint()
 										break
 									}
 								}
 								// now execute the tool like normal
 								const content = await extractTextFromFile(absolutePath)
 								pushToolResult(content)
-								await this.saveCheckpoint()
+
 								break
 							}
 						} catch (error) {
 							await handleError("reading file", error)
-							await this.saveCheckpoint()
+
 							break
 						}
 					}
@@ -1947,7 +1987,7 @@ export class Cline {
 								if (!relDirPath) {
 									this.consecutiveMistakeCount++
 									pushToolResult(await this.sayAndCreateMissingParamError("list_files", "path"))
-									await this.saveCheckpoint()
+
 									break
 								}
 								this.consecutiveMistakeCount = 0
@@ -1977,17 +2017,16 @@ export class Cline {
 									this.removeLastPartialMessageIfExistsWithType("say", "tool")
 									const didApprove = await askApproval("tool", completeMessage)
 									if (!didApprove) {
-										await this.saveCheckpoint()
 										break
 									}
 								}
 								pushToolResult(result)
-								await this.saveCheckpoint()
+
 								break
 							}
 						} catch (error) {
 							await handleError("listing files", error)
-							await this.saveCheckpoint()
+
 							break
 						}
 					}
@@ -2015,7 +2054,7 @@ export class Cline {
 								if (!relDirPath) {
 									this.consecutiveMistakeCount++
 									pushToolResult(await this.sayAndCreateMissingParamError("list_code_definition_names", "path"))
-									await this.saveCheckpoint()
+
 									break
 								}
 
@@ -2042,17 +2081,16 @@ export class Cline {
 									this.removeLastPartialMessageIfExistsWithType("say", "tool")
 									const didApprove = await askApproval("tool", completeMessage)
 									if (!didApprove) {
-										await this.saveCheckpoint()
 										break
 									}
 								}
 								pushToolResult(result)
-								await this.saveCheckpoint()
+
 								break
 							}
 						} catch (error) {
 							await handleError("parsing source code definitions", error)
-							await this.saveCheckpoint()
+
 							break
 						}
 					}
@@ -2084,13 +2122,13 @@ export class Cline {
 								if (!relDirPath) {
 									this.consecutiveMistakeCount++
 									pushToolResult(await this.sayAndCreateMissingParamError("search_files", "path"))
-									await this.saveCheckpoint()
+
 									break
 								}
 								if (!regex) {
 									this.consecutiveMistakeCount++
 									pushToolResult(await this.sayAndCreateMissingParamError("search_files", "regex"))
-									await this.saveCheckpoint()
+
 									break
 								}
 								this.consecutiveMistakeCount = 0
@@ -2119,17 +2157,16 @@ export class Cline {
 									this.removeLastPartialMessageIfExistsWithType("say", "tool")
 									const didApprove = await askApproval("tool", completeMessage)
 									if (!didApprove) {
-										await this.saveCheckpoint()
 										break
 									}
 								}
 								pushToolResult(results)
-								await this.saveCheckpoint()
+
 								break
 							}
 						} catch (error) {
 							await handleError("searching files", error)
-							await this.saveCheckpoint()
+
 							break
 						}
 					}
@@ -2188,7 +2225,7 @@ export class Cline {
 										this.consecutiveMistakeCount++
 										pushToolResult(await this.sayAndCreateMissingParamError("browser_action", "url"))
 										await this.browserSession.closeBrowser()
-										await this.saveCheckpoint()
+
 										break
 									}
 									this.consecutiveMistakeCount = 0
@@ -2204,7 +2241,6 @@ export class Cline {
 										this.removeLastPartialMessageIfExistsWithType("say", "browser_action_launch")
 										const didApprove = await askApproval("browser_action_launch", url)
 										if (!didApprove) {
-											await this.saveCheckpoint()
 											break
 										}
 									}
@@ -2223,7 +2259,7 @@ export class Cline {
 												await this.sayAndCreateMissingParamError("browser_action", "coordinate"),
 											)
 											await this.browserSession.closeBrowser()
-											await this.saveCheckpoint()
+
 											break // can't be within an inner switch
 										}
 									}
@@ -2232,7 +2268,7 @@ export class Cline {
 											this.consecutiveMistakeCount++
 											pushToolResult(await this.sayAndCreateMissingParamError("browser_action", "text"))
 											await this.browserSession.closeBrowser()
-											await this.saveCheckpoint()
+
 											break
 										}
 									}
@@ -2281,7 +2317,7 @@ export class Cline {
 												browserActionResult.screenshot ? [browserActionResult.screenshot] : [],
 											),
 										)
-										await this.saveCheckpoint()
+
 										break
 									case "close":
 										pushToolResult(
@@ -2289,17 +2325,16 @@ export class Cline {
 												`The browser has been closed. You may now proceed to using other tools.`,
 											),
 										)
-										await this.saveCheckpoint()
+
 										break
 								}
 
-								await this.saveCheckpoint()
 								break
 							}
 						} catch (error) {
 							await this.browserSession.closeBrowser() // if any error occurs, the browser session is terminated
 							await handleError("executing browser action", error)
-							await this.saveCheckpoint()
+
 							break
 						}
 					}
@@ -2327,7 +2362,7 @@ export class Cline {
 								if (!command) {
 									this.consecutiveMistakeCount++
 									pushToolResult(await this.sayAndCreateMissingParamError("execute_command", "command"))
-									await this.saveCheckpoint()
+
 									break
 								}
 								if (!requiresApprovalRaw) {
@@ -2335,7 +2370,7 @@ export class Cline {
 									pushToolResult(
 										await this.sayAndCreateMissingParamError("execute_command", "requires_approval"),
 									)
-									await this.saveCheckpoint()
+
 									break
 								}
 								this.consecutiveMistakeCount = 0
@@ -2346,7 +2381,7 @@ export class Cline {
 									pushToolResult(
 										formatResponse.toolError(formatResponse.clineIgnoreError(ignoredFileAttemptedToAccess)),
 									)
-									await this.saveCheckpoint()
+
 									break
 								}
 
@@ -2368,7 +2403,6 @@ export class Cline {
 											`${this.shouldAutoApproveTool(block.name) && requiresApproval ? COMMAND_REQ_APP_STRING : ""}`, // ugly hack until we refactor combineCommandSequences
 									)
 									if (!didApprove) {
-										await this.saveCheckpoint()
 										break
 									}
 								}
@@ -2397,12 +2431,14 @@ export class Cline {
 								this.providerRef.deref()?.workspaceTracker?.populateFilePaths()
 
 								pushToolResult(result)
+
 								await this.saveCheckpoint()
+
 								break
 							}
 						} catch (error) {
 							await handleError("executing command", error)
-							await this.saveCheckpoint()
+
 							break
 						}
 					}
@@ -2432,13 +2468,13 @@ export class Cline {
 								if (!server_name) {
 									this.consecutiveMistakeCount++
 									pushToolResult(await this.sayAndCreateMissingParamError("use_mcp_tool", "server_name"))
-									await this.saveCheckpoint()
+
 									break
 								}
 								if (!tool_name) {
 									this.consecutiveMistakeCount++
 									pushToolResult(await this.sayAndCreateMissingParamError("use_mcp_tool", "tool_name"))
-									await this.saveCheckpoint()
+
 									break
 								}
 								// arguments are optional, but if they are provided they must be valid JSON
@@ -2462,7 +2498,7 @@ export class Cline {
 												formatResponse.invalidMcpToolArgumentError(server_name, tool_name),
 											),
 										)
-										await this.saveCheckpoint()
+
 										break
 									}
 								}
@@ -2490,7 +2526,6 @@ export class Cline {
 									this.removeLastPartialMessageIfExistsWithType("say", "use_mcp_server")
 									const didApprove = await askApproval("use_mcp_server", completeMessage)
 									if (!didApprove) {
-										await this.saveCheckpoint()
 										break
 									}
 								}
@@ -2519,12 +2554,14 @@ export class Cline {
 											.join("\n\n") || "(No response)"
 								await this.say("mcp_server_response", toolResultPretty)
 								pushToolResult(formatResponse.toolResult(toolResultPretty))
+
 								await this.saveCheckpoint()
+
 								break
 							}
 						} catch (error) {
 							await handleError("executing MCP tool", error)
-							await this.saveCheckpoint()
+
 							break
 						}
 					}
@@ -2552,13 +2589,13 @@ export class Cline {
 								if (!server_name) {
 									this.consecutiveMistakeCount++
 									pushToolResult(await this.sayAndCreateMissingParamError("access_mcp_resource", "server_name"))
-									await this.saveCheckpoint()
+
 									break
 								}
 								if (!uri) {
 									this.consecutiveMistakeCount++
 									pushToolResult(await this.sayAndCreateMissingParamError("access_mcp_resource", "uri"))
-									await this.saveCheckpoint()
+
 									break
 								}
 								this.consecutiveMistakeCount = 0
@@ -2579,7 +2616,6 @@ export class Cline {
 									this.removeLastPartialMessageIfExistsWithType("say", "use_mcp_server")
 									const didApprove = await askApproval("use_mcp_server", completeMessage)
 									if (!didApprove) {
-										await this.saveCheckpoint()
 										break
 									}
 								}
@@ -2599,12 +2635,12 @@ export class Cline {
 										.join("\n\n") || "(Empty response)"
 								await this.say("mcp_server_response", resourceResultPretty)
 								pushToolResult(formatResponse.toolResult(resourceResultPretty))
-								await this.saveCheckpoint()
+
 								break
 							}
 						} catch (error) {
 							await handleError("accessing MCP resource", error)
-							await this.saveCheckpoint()
+
 							break
 						}
 					}
@@ -2618,7 +2654,7 @@ export class Cline {
 								if (!question) {
 									this.consecutiveMistakeCount++
 									pushToolResult(await this.sayAndCreateMissingParamError("ask_followup_question", "question"))
-									await this.saveCheckpoint()
+
 									break
 								}
 								this.consecutiveMistakeCount = 0
@@ -2633,12 +2669,12 @@ export class Cline {
 								const { text, images } = await this.ask("followup", question, false)
 								await this.say("user_feedback", text ?? "", images)
 								pushToolResult(formatResponse.toolResult(`<answer>\n${text}\n</answer>`, images))
-								await this.saveCheckpoint()
+
 								break
 							}
 						} catch (error) {
 							await handleError("asking question", error)
-							await this.saveCheckpoint()
+
 							break
 						}
 					}
@@ -2654,7 +2690,7 @@ export class Cline {
 								if (!response) {
 									this.consecutiveMistakeCount++
 									pushToolResult(await this.sayAndCreateMissingParamError("plan_mode_response", "response"))
-									// await this.saveCheckpoint()
+									//
 									break
 								}
 								this.consecutiveMistakeCount = 0
@@ -2683,12 +2719,12 @@ export class Cline {
 									pushToolResult(formatResponse.toolResult(`<user_message>\n${text}\n</user_message>`, images))
 								}
 
-								// await this.saveCheckpoint()
+								//
 								break
 							}
 						} catch (error) {
 							await handleError("responding to inquiry", error)
-							// await this.saveCheckpoint()
+							//
 							break
 						}
 					}
@@ -2749,7 +2785,7 @@ export class Cline {
 										// last message is completion_result
 										// we have command string, which means we have the result as well, so finish it (doesnt have to exist yet)
 										await this.say("completion_result", removeClosingTag("result", result), undefined, false)
-										await this.saveCheckpoint()
+										await this.saveCheckpoint(true)
 										await addNewChangesFlagToLastCompletionResultMessage()
 										await this.ask("command", removeClosingTag("command", command), block.partial).catch(
 											() => {},
@@ -2769,7 +2805,6 @@ export class Cline {
 								if (!result) {
 									this.consecutiveMistakeCount++
 									pushToolResult(await this.sayAndCreateMissingParamError("attempt_completion", "result"))
-									await this.saveCheckpoint()
 									break
 								}
 								this.consecutiveMistakeCount = 0
@@ -2786,31 +2821,29 @@ export class Cline {
 									if (lastMessage && lastMessage.ask !== "command") {
 										// havent sent a command message yet so first send completion_result then command
 										await this.say("completion_result", result, undefined, false)
-										await this.saveCheckpoint()
+										await this.saveCheckpoint(true)
 										await addNewChangesFlagToLastCompletionResultMessage()
 									} else {
 										// we already sent a command message, meaning the complete completion message has also been sent
-										await this.saveCheckpoint()
+										await this.saveCheckpoint(true)
 									}
 
 									// complete command message
 									const didApprove = await askApproval("command", command)
 									if (!didApprove) {
-										await this.saveCheckpoint()
 										break
 									}
 									const [userRejected, execCommandResult] = await this.executeCommandTool(command!)
 									if (userRejected) {
 										this.didRejectTool = true
 										pushToolResult(execCommandResult)
-										await this.saveCheckpoint()
 										break
 									}
 									// user didn't reject, but the command may have output
 									commandResult = execCommandResult
 								} else {
 									await this.say("completion_result", result, undefined, false)
-									await this.saveCheckpoint()
+									await this.saveCheckpoint(true)
 									await addNewChangesFlagToLastCompletionResultMessage()
 								}
 
@@ -2844,12 +2877,12 @@ export class Cline {
 								})
 								this.userMessageContent.push(...toolResults)
 
-								// await this.saveCheckpoint()
+								//
 								break
 							}
 						} catch (error) {
 							await handleError("attempting completion", error)
-							await this.saveCheckpoint()
+
 							break
 						}
 					}
diff --git a/src/shared/ExtensionMessage.ts b/src/shared/ExtensionMessage.ts
index c08d611eda..70a58fbfd8 100644
--- a/src/shared/ExtensionMessage.ts
+++ b/src/shared/ExtensionMessage.ts
@@ -84,6 +84,7 @@ export interface ClineMessage {
 	images?: string[]
 	partial?: boolean
 	lastCheckpointHash?: string
+	isCheckpointCheckedOut?: boolean
 	conversationHistoryIndex?: number
 	conversationHistoryDeletedRange?: [number, number] // for when conversation history is truncated for API requests
 }
@@ -127,6 +128,7 @@ export type ClineSay =
 	| "diff_error"
 	| "deleted_api_reqs"
 	| "clineignore_error"
+	| "checkpoint_created"
 
 export interface ClineSayTool {
 	tool:
diff --git a/webview-ui/package-lock.json b/webview-ui/package-lock.json
index 5ac343c886..4a06dc2e45 100644
--- a/webview-ui/package-lock.json
+++ b/webview-ui/package-lock.json
@@ -8,6 +8,7 @@
 			"name": "webview-ui",
 			"version": "0.1.0",
 			"dependencies": {
+				"@floating-ui/react": "^0.27.4",
 				"@vscode/webview-ui-toolkit": "^1.4.0",
 				"debounce": "^2.1.1",
 				"fast-deep-equal": "^3.1.3",
@@ -3001,6 +3002,65 @@
 				"node": "^12.22.0 || ^14.17.0 || >=16.0.0"
 			}
 		},
+		"node_modules/@floating-ui/core": {
+			"version": "1.6.9",
+			"resolved": "https://registry.npmjs.org/@floating-ui/core/-/core-1.6.9.tgz",
+			"integrity": "sha512-uMXCuQ3BItDUbAMhIXw7UPXRfAlOAvZzdK9BWpE60MCn+Svt3aLn9jsPTi/WNGlRUu2uI0v5S7JiIUsbsvh3fw==",
+			"license": "MIT",
+			"dependencies": {
+				"@floating-ui/utils": "^0.2.9"
+			}
+		},
+		"node_modules/@floating-ui/dom": {
+			"version": "1.6.13",
+			"resolved": "https://registry.npmjs.org/@floating-ui/dom/-/dom-1.6.13.tgz",
+			"integrity": "sha512-umqzocjDgNRGTuO7Q8CU32dkHkECqI8ZdMZ5Swb6QAM0t5rnlrN3lGo1hdpscRd3WS8T6DKYK4ephgIH9iRh3w==",
+			"license": "MIT",
+			"dependencies": {
+				"@floating-ui/core": "^1.6.0",
+				"@floating-ui/utils": "^0.2.9"
+			}
+		},
+		"node_modules/@floating-ui/react": {
+			"version": "0.27.4",
+			"resolved": "https://registry.npmjs.org/@floating-ui/react/-/react-0.27.4.tgz",
+			"integrity": "sha512-05mXdkUiVh8NCEcYKQ2C9SV9IkZ9k/dFtYmaEIN2riLv80UHoXylgBM76cgPJYfLJM3dJz7UE5MOVH0FypMd2Q==",
+			"license": "MIT",
+			"dependencies": {
+				"@floating-ui/react-dom": "^2.1.2",
+				"@floating-ui/utils": "^0.2.9",
+				"tabbable": "^6.0.0"
+			},
+			"peerDependencies": {
+				"react": ">=17.0.0",
+				"react-dom": ">=17.0.0"
+			}
+		},
+		"node_modules/@floating-ui/react-dom": {
+			"version": "2.1.2",
+			"resolved": "https://registry.npmjs.org/@floating-ui/react-dom/-/react-dom-2.1.2.tgz",
+			"integrity": "sha512-06okr5cgPzMNBy+Ycse2A6udMi4bqwW/zgBF/rwjcNqWkyr82Mcg8b0vjX8OJpZFy/FKjJmw6wV7t44kK6kW7A==",
+			"license": "MIT",
+			"dependencies": {
+				"@floating-ui/dom": "^1.0.0"
+			},
+			"peerDependencies": {
+				"react": ">=16.8.0",
+				"react-dom": ">=16.8.0"
+			}
+		},
+		"node_modules/@floating-ui/react/node_modules/tabbable": {
+			"version": "6.2.0",
+			"resolved": "https://registry.npmjs.org/tabbable/-/tabbable-6.2.0.tgz",
+			"integrity": "sha512-Cat63mxsVJlzYvN51JmVXIgNoUokrIaT2zLclCXjRd8boZ0004U4KCs/sToJ75C6sdlByWxpYnb5Boif1VSFew==",
+			"license": "MIT"
+		},
+		"node_modules/@floating-ui/utils": {
+			"version": "0.2.9",
+			"resolved": "https://registry.npmjs.org/@floating-ui/utils/-/utils-0.2.9.tgz",
+			"integrity": "sha512-MDWhGtE+eHw5JW7lq4qhc5yRLS11ERl1c7Z6Xd0a58DozHES6EnNNwUWbMiG4J9Cgj053Bhk8zvlhFYKVhULwg==",
+			"license": "MIT"
+		},
 		"node_modules/@humanwhocodes/config-array": {
 			"version": "0.13.0",
 			"resolved": "https://registry.npmjs.org/@humanwhocodes/config-array/-/config-array-0.13.0.tgz",
diff --git a/webview-ui/package.json b/webview-ui/package.json
index 8d89ea490f..2f023747d1 100644
--- a/webview-ui/package.json
+++ b/webview-ui/package.json
@@ -3,6 +3,7 @@
 	"version": "0.1.0",
 	"private": true,
 	"dependencies": {
+		"@floating-ui/react": "^0.27.4",
 		"@vscode/webview-ui-toolkit": "^1.4.0",
 		"debounce": "^2.1.1",
 		"fast-deep-equal": "^3.1.3",
@@ -53,11 +54,11 @@
 		"@testing-library/jest-dom": "^5.17.0",
 		"@testing-library/react": "^15.0.6",
 		"@testing-library/user-event": "^13.5.0",
-		"@types/vscode-webview": "^1.57.5",
 		"@types/jest": "^27.5.2",
 		"@types/node": "^20.x",
 		"@types/react": "^18.3.3",
 		"@types/react-dom": "^18.3.0",
+		"@types/vscode-webview": "^1.57.5",
 		"jsdom": "^25.0.1",
 		"vitest": "^2.1.9"
 	}
diff --git a/webview-ui/src/components/chat/BrowserSessionRow.tsx b/webview-ui/src/components/chat/BrowserSessionRow.tsx
index 7153cb7f21..e79f4be3f7 100644
--- a/webview-ui/src/components/chat/BrowserSessionRow.tsx
+++ b/webview-ui/src/components/chat/BrowserSessionRow.tsx
@@ -1,17 +1,16 @@
+import { VSCodeButton } from "@vscode/webview-ui-toolkit/react"
 import deepEqual from "fast-deep-equal"
 import React, { memo, useEffect, useMemo, useRef, useState } from "react"
 import { useSize } from "react-use"
+import styled from "styled-components"
+import { BROWSER_VIEWPORT_PRESETS } from "../../../../src/shared/BrowserSettings"
 import { BrowserAction, BrowserActionResult, ClineMessage, ClineSayBrowserAction } from "../../../../src/shared/ExtensionMessage"
+import { useExtensionState } from "../../context/ExtensionStateContext"
 import { vscode } from "../../utils/vscode"
+import { BrowserSettingsMenu } from "../browser/BrowserSettingsMenu"
+import { CheckpointControls } from "../common/CheckpointControls"
 import CodeBlock, { CODE_BLOCK_BG_COLOR } from "../common/CodeBlock"
 import { ChatRowContent, ProgressIndicator } from "./ChatRow"
-import { VSCodeButton } from "@vscode/webview-ui-toolkit/react"
-import styled from "styled-components"
-import { CheckpointControls, CheckpointOverlay } from "../common/CheckpointControls"
-import { findLast } from "../../../../src/shared/array"
-import { BrowserSettingsMenu } from "../browser/BrowserSettingsMenu"
-import { useExtensionState } from "../../context/ExtensionStateContext"
-import { BROWSER_VIEWPORT_PRESETS } from "../../../../src/shared/BrowserSettings"
 
 interface BrowserSessionRowProps {
 	messages: ClineMessage[]
@@ -144,10 +143,10 @@ const BrowserSessionRow = memo((props: BrowserSessionRowProps) => {
 		return launchMessage?.say === "browser_action_launch"
 	}, [messages])
 
-	const lastCheckpointMessageTs = useMemo(() => {
-		const lastCheckpointMessage = findLast(messages, (m) => m.lastCheckpointHash !== undefined)
-		return lastCheckpointMessage?.ts
-	}, [messages])
+	// const lastCheckpointMessageTs = useMemo(() => {
+	// 	const lastCheckpointMessage = findLast(messages, (m) => m.lastCheckpointHash !== undefined)
+	// 	return lastCheckpointMessage?.ts
+	// }, [messages])
 
 	// Find the latest available URL and screenshot
 	const latestState = useMemo(() => {
@@ -231,10 +230,10 @@ const BrowserSessionRow = memo((props: BrowserSessionRowProps) => {
 	// Use latest click position while browsing, otherwise use display state
 	const mousePosition = isBrowsing ? latestClickPosition || displayState.mousePosition : displayState.mousePosition
 
-	let shouldShowCheckpoints = true
-	if (isLast) {
-		shouldShowCheckpoints = lastModifiedMessage?.ask === "resume_completed_task" || lastModifiedMessage?.ask === "resume_task"
-	}
+	// let shouldShowCheckpoints = true
+	// if (isLast) {
+	// 	shouldShowCheckpoints = lastModifiedMessage?.ask === "resume_completed_task" || lastModifiedMessage?.ask === "resume_task"
+	// }
 
 	const shouldShowSettings = useMemo(() => {
 		const lastMessage = messages[messages.length - 1]
@@ -423,7 +422,7 @@ const BrowserSessionRow = memo((props: BrowserSessionRowProps) => {
 				</div>
 			)}
 
-			{shouldShowCheckpoints && <CheckpointOverlay messageTs={lastCheckpointMessageTs} />}
+			{/* {shouldShowCheckpoints && <CheckpointOverlay messageTs={lastCheckpointMessageTs} />} */}
 		</BrowserSessionRowContainer>,
 	)
 
diff --git a/webview-ui/src/components/chat/ChatRow.tsx b/webview-ui/src/components/chat/ChatRow.tsx
index e4fc6a0aa7..889ba0efb7 100644
--- a/webview-ui/src/components/chat/ChatRow.tsx
+++ b/webview-ui/src/components/chat/ChatRow.tsx
@@ -24,6 +24,7 @@ import Thumbnails from "../common/Thumbnails"
 import McpResourceRow from "../mcp/McpResourceRow"
 import McpToolRow from "../mcp/McpToolRow"
 import { highlightMentions } from "./TaskHeader"
+import { CheckmarkControl } from "../common/CheckmarkControl"
 
 const ChatRowContainer = styled.div`
 	padding: 10px 6px 10px 15px;
@@ -59,8 +60,8 @@ const ChatRow = memo(
 				message.ask === "tool" ||
 				message.say === "command" ||
 				message.ask === "command" ||
-				message.say === "completion_result" ||
-				message.ask === "completion_result" ||
+				// message.say === "completion_result" ||
+				// message.ask === "completion_result" ||
 				message.say === "use_mcp_server" ||
 				message.ask === "use_mcp_server")
 
@@ -999,6 +1000,12 @@ export const ChatRowContent = ({ message, isExpanded, onToggleExpand, lastModifi
 							</div>
 						</>
 					)
+				case "checkpoint_created":
+					return (
+						<>
+							<CheckmarkControl messageTs={message.ts} isCheckpointCheckedOut={message.isCheckpointCheckedOut} />
+						</>
+					)
 				case "completion_result":
 					const hasChanges = message.text?.endsWith(COMPLETION_RESULT_CHANGES_FLAG) ?? false
 					const text = hasChanges ? message.text?.slice(0, -COMPLETION_RESULT_CHANGES_FLAG.length) : message.text
diff --git a/webview-ui/src/components/chat/TaskHeader.tsx b/webview-ui/src/components/chat/TaskHeader.tsx
index 92e44e2e4d..f6b245d5f6 100644
--- a/webview-ui/src/components/chat/TaskHeader.tsx
+++ b/webview-ui/src/components/chat/TaskHeader.tsx
@@ -422,7 +422,7 @@ const TaskHeader: React.FC<TaskHeaderProps> = ({
 									</span>
 								</div>
 							)}
-							{ContextWindowComponent}
+							{/* {ContextWindowComponent} */}
 							{isCostAvailable && (
 								<div
 									style={{
diff --git a/webview-ui/src/components/common/CheckmarkControl.tsx b/webview-ui/src/components/common/CheckmarkControl.tsx
new file mode 100644
index 0000000000..31043fb09f
--- /dev/null
+++ b/webview-ui/src/components/common/CheckmarkControl.tsx
@@ -0,0 +1,394 @@
+import { useCallback, useRef, useState, useEffect } from "react"
+import { useClickAway, useEvent } from "react-use"
+import styled from "styled-components"
+import { ExtensionMessage } from "../../../../src/shared/ExtensionMessage"
+import { vscode } from "../../utils/vscode"
+import { CODE_BLOCK_BG_COLOR } from "./CodeBlock"
+import { ClineCheckpointRestore } from "../../../../src/shared/WebviewMessage"
+import { VSCodeButton } from "@vscode/webview-ui-toolkit/react"
+import { createPortal } from "react-dom"
+import { useFloating, offset, flip, shift } from "@floating-ui/react"
+
+interface CheckmarkControlProps {
+	messageTs?: number
+	isCheckpointCheckedOut?: boolean
+}
+
+export const CheckmarkControl = ({ messageTs, isCheckpointCheckedOut }: CheckmarkControlProps) => {
+	const [compareDisabled, setCompareDisabled] = useState(false)
+	const [restoreTaskDisabled, setRestoreTaskDisabled] = useState(false)
+	const [restoreWorkspaceDisabled, setRestoreWorkspaceDisabled] = useState(false)
+	const [restoreBothDisabled, setRestoreBothDisabled] = useState(false)
+	const [showRestoreConfirm, setShowRestoreConfirm] = useState(false)
+	const [hasMouseEntered, setHasMouseEntered] = useState(false)
+	const containerRef = useRef<HTMLDivElement>(null)
+	const tooltipRef = useRef<HTMLDivElement>(null)
+
+	const { refs, floatingStyles, update, placement } = useFloating({
+		placement: "bottom-end",
+		middleware: [
+			offset({
+				mainAxis: 8,
+				crossAxis: 10,
+			}),
+			flip(),
+			shift(),
+		],
+	})
+
+	useEffect(() => {
+		const handleScroll = () => {
+			update()
+		}
+		window.addEventListener("scroll", handleScroll, true)
+		return () => window.removeEventListener("scroll", handleScroll, true)
+	}, [update])
+
+	useEffect(() => {
+		if (showRestoreConfirm) {
+			update()
+		}
+	}, [showRestoreConfirm, update])
+
+	const handleMessage = useCallback((event: MessageEvent<ExtensionMessage>) => {
+		if (event.data.type === "relinquishControl") {
+			setCompareDisabled(false)
+			setRestoreTaskDisabled(false)
+			setRestoreWorkspaceDisabled(false)
+			setRestoreBothDisabled(false)
+			setShowRestoreConfirm(false)
+		}
+	}, [])
+
+	const handleRestoreTask = () => {
+		setRestoreTaskDisabled(true)
+		vscode.postMessage({
+			type: "checkpointRestore",
+			number: messageTs,
+			text: "task",
+		})
+	}
+
+	const handleRestoreWorkspace = () => {
+		setRestoreWorkspaceDisabled(true)
+		vscode.postMessage({
+			type: "checkpointRestore",
+			number: messageTs,
+			text: "workspace",
+		})
+	}
+
+	const handleRestoreBoth = () => {
+		setRestoreBothDisabled(true)
+		vscode.postMessage({
+			type: "checkpointRestore",
+			number: messageTs,
+			text: "taskAndWorkspace",
+		})
+	}
+
+	const handleMouseEnter = () => {
+		setHasMouseEntered(true)
+	}
+
+	const handleMouseLeave = () => {
+		if (hasMouseEntered) {
+			setShowRestoreConfirm(false)
+			setHasMouseEntered(false)
+		}
+	}
+
+	const handleControlsMouseLeave = (e: React.MouseEvent) => {
+		const tooltipElement = tooltipRef.current
+
+		if (tooltipElement && showRestoreConfirm) {
+			const tooltipRect = tooltipElement.getBoundingClientRect()
+
+			if (
+				e.clientY >= tooltipRect.top &&
+				e.clientY <= tooltipRect.bottom &&
+				e.clientX >= tooltipRect.left &&
+				e.clientX <= tooltipRect.right
+			) {
+				return
+			}
+		}
+
+		setShowRestoreConfirm(false)
+		setHasMouseEntered(false)
+	}
+
+	useEvent("message", handleMessage)
+
+	return (
+		<Container isMenuOpen={showRestoreConfirm} $isCheckedOut={isCheckpointCheckedOut} onMouseLeave={handleControlsMouseLeave}>
+			<i
+				className="codicon codicon-bookmark"
+				style={{
+					color: isCheckpointCheckedOut ? "var(--vscode-textLink-foreground)" : "var(--vscode-descriptionForeground)",
+					fontSize: "12px",
+					flexShrink: 0,
+				}}
+			/>
+			<Label $isCheckedOut={isCheckpointCheckedOut}>
+				{isCheckpointCheckedOut ? "Checkpoint (restored)" : "Checkpoint"}
+			</Label>
+			<DottedLine $isCheckedOut={isCheckpointCheckedOut} />
+			<ButtonGroup>
+				<CustomButton
+					$isCheckedOut={isCheckpointCheckedOut}
+					disabled={compareDisabled}
+					style={{ cursor: compareDisabled ? "wait" : "pointer" }}
+					onClick={() => {
+						setCompareDisabled(true)
+						vscode.postMessage({
+							type: "checkpointDiff",
+							number: messageTs,
+						})
+					}}>
+					Compare
+				</CustomButton>
+				<DottedLine small $isCheckedOut={isCheckpointCheckedOut} />
+				<div ref={refs.setReference} style={{ position: "relative" }}>
+					<CustomButton
+						$isCheckedOut={isCheckpointCheckedOut}
+						isActive={showRestoreConfirm}
+						onClick={() => setShowRestoreConfirm(true)}>
+						Restore
+					</CustomButton>
+					{showRestoreConfirm &&
+						createPortal(
+							<RestoreConfirmTooltip
+								ref={refs.setFloating}
+								style={floatingStyles}
+								data-placement={placement}
+								onMouseEnter={handleMouseEnter}
+								onMouseLeave={handleMouseLeave}>
+								<RestoreOption>
+									<VSCodeButton
+										onClick={handleRestoreWorkspace}
+										disabled={restoreWorkspaceDisabled}
+										style={{
+											cursor: restoreWorkspaceDisabled ? "wait" : "pointer",
+											width: "100%",
+											marginBottom: "10px",
+										}}>
+										Restore Files
+									</VSCodeButton>
+									<p>
+										Restores your project's files back to a snapshot taken at this point (use "Compare" to see
+										what will be reverted)
+									</p>
+								</RestoreOption>
+								{/* <RestoreOption>
+									<VSCodeButton
+										onClick={handleRestoreTask}
+										disabled={restoreTaskDisabled}
+										style={{
+											cursor: restoreTaskDisabled ? "wait" : "pointer",
+											width: "100%",
+											marginBottom: "10px",
+										}}>
+										Restore Task Only
+									</VSCodeButton>
+									<p>Deletes messages after this point (does not affect workspace files)</p>
+								</RestoreOption> */}
+								<RestoreOption>
+									<VSCodeButton
+										onClick={handleRestoreBoth}
+										disabled={restoreBothDisabled}
+										style={{
+											cursor: restoreBothDisabled ? "wait" : "pointer",
+											width: "100%",
+											marginBottom: "10px",
+										}}>
+										Restore Files & Task
+									</VSCodeButton>
+									<p>Restores your project's files and deletes all messages after this point</p>
+								</RestoreOption>
+							</RestoreConfirmTooltip>,
+							document.body,
+						)}
+				</div>
+				<DottedLine small $isCheckedOut={isCheckpointCheckedOut} />
+			</ButtonGroup>
+		</Container>
+	)
+}
+
+const Container = styled.div<{ isMenuOpen?: boolean; $isCheckedOut?: boolean }>`
+	display: flex;
+	align-items: center;
+	padding: 4px 0;
+	gap: 4px;
+	position: relative;
+	min-width: 0;
+	margin-top: -10px;
+	margin-bottom: -10px;
+	opacity: ${(props) => (props.$isCheckedOut ? 1 : props.isMenuOpen ? 1 : 0.5)};
+
+	&:hover {
+		opacity: 1;
+	}
+`
+
+const Label = styled.span<{ $isCheckedOut?: boolean }>`
+	color: ${(props) => (props.$isCheckedOut ? "var(--vscode-textLink-foreground)" : "var(--vscode-descriptionForeground)")};
+	font-size: 9px;
+	flex-shrink: 0;
+`
+
+const DottedLine = styled.div<{ small?: boolean; $isCheckedOut?: boolean }>`
+	flex: ${(props) => (props.small ? "0 0 5px" : "1")};
+	min-width: ${(props) => (props.small ? "5px" : "5px")};
+	height: 1px;
+	background-image: linear-gradient(
+		to right,
+		${(props) => (props.$isCheckedOut ? "var(--vscode-textLink-foreground)" : "var(--vscode-descriptionForeground)")} 50%,
+		transparent 50%
+	);
+	background-size: 4px 1px;
+	background-repeat: repeat-x;
+`
+
+const ButtonGroup = styled.div`
+	display: flex;
+	align-items: center;
+	gap: 4px;
+	flex-shrink: 0;
+`
+
+const CustomButton = styled.button<{ disabled?: boolean; isActive?: boolean; $isCheckedOut?: boolean }>`
+	background: ${(props) =>
+		props.isActive || props.disabled
+			? props.$isCheckedOut
+				? "var(--vscode-textLink-foreground)"
+				: "var(--vscode-descriptionForeground)"
+			: "transparent"};
+	border: none;
+	color: ${(props) =>
+		props.isActive || props.disabled
+			? "var(--vscode-editor-background)"
+			: props.$isCheckedOut
+				? "var(--vscode-textLink-foreground)"
+				: "var(--vscode-descriptionForeground)"};
+	padding: 2px 6px;
+	font-size: 9px;
+	cursor: pointer;
+	position: relative;
+
+	&::before {
+		content: "";
+		position: absolute;
+		top: 0;
+		left: 0;
+		right: 0;
+		bottom: 0;
+		border-radius: 1px;
+		background-image: ${(props) =>
+			props.isActive || props.disabled
+				? "none"
+				: `linear-gradient(to right, ${props.$isCheckedOut ? "var(--vscode-textLink-foreground)" : "var(--vscode-descriptionForeground)"} 50%, transparent 50%),
+			linear-gradient(to bottom, ${props.$isCheckedOut ? "var(--vscode-textLink-foreground)" : "var(--vscode-descriptionForeground)"} 50%, transparent 50%),
+			linear-gradient(to right, ${props.$isCheckedOut ? "var(--vscode-textLink-foreground)" : "var(--vscode-descriptionForeground)"} 50%, transparent 50%),
+			linear-gradient(to bottom, ${props.$isCheckedOut ? "var(--vscode-textLink-foreground)" : "var(--vscode-descriptionForeground)"} 50%, transparent 50%)`};
+		background-size: ${(props) => (props.isActive || props.disabled ? "auto" : `4px 1px, 1px 4px, 4px 1px, 1px 4px`)};
+		background-repeat: repeat-x, repeat-y, repeat-x, repeat-y;
+		background-position:
+			0 0,
+			100% 0,
+			0 100%,
+			0 0;
+	}
+
+	&:hover:not(:disabled) {
+		background: ${(props) =>
+			props.$isCheckedOut ? "var(--vscode-textLink-foreground)" : "var(--vscode-descriptionForeground)"};
+		color: var(--vscode-editor-background);
+		&::before {
+			display: none;
+		}
+	}
+
+	&:disabled {
+		opacity: 0.5;
+		cursor: not-allowed;
+	}
+`
+
+const RestoreOption = styled.div`
+	&:not(:last-child) {
+		margin-bottom: 10px;
+		padding-bottom: 4px;
+		border-bottom: 1px solid var(--vscode-editorGroup-border);
+	}
+
+	p {
+		margin: 0 0 2px 0;
+		color: var(--vscode-descriptionForeground);
+		font-size: 11px;
+		line-height: 14px;
+	}
+
+	&:last-child p {
+		margin: 0 0 -2px 0;
+	}
+`
+
+const RestoreConfirmTooltip = styled.div`
+	position: fixed;
+	background: ${CODE_BLOCK_BG_COLOR};
+	border: 1px solid var(--vscode-editorGroup-border);
+	padding: 12px;
+	border-radius: 3px;
+	width: min(calc(100vw - 54px), 600px);
+	z-index: 1000;
+
+	// Add invisible padding to create a safe hover zone
+	&::before {
+		content: "";
+		position: absolute;
+		top: -8px;
+		left: 0;
+		right: 0;
+		height: 8px;
+	}
+
+	// Adjust arrow to be above the padding
+	&::after {
+		content: "";
+		position: absolute;
+		top: -6px;
+		right: 24px;
+		width: 10px;
+		height: 10px;
+		background: ${CODE_BLOCK_BG_COLOR};
+		border-left: 1px solid var(--vscode-editorGroup-border);
+		border-top: 1px solid var(--vscode-editorGroup-border);
+		transform: rotate(45deg);
+		z-index: 1;
+	}
+
+	// When menu appears above the button
+	&[data-placement^="top"] {
+		&::before {
+			top: auto;
+			bottom: -8px;
+		}
+
+		&::after {
+			top: auto;
+			bottom: -6px;
+			right: 24px;
+			transform: rotate(225deg);
+		}
+	}
+
+	p {
+		margin: 0 0 6px 0;
+		color: var(--vscode-descriptionForeground);
+		font-size: 12px;
+		white-space: normal;
+		word-wrap: break-word;
+	}
+`