From 1a527fca123db20c321a81e2daef90c13f5703a0 Mon Sep 17 00:00:00 2001
From: cte <cestreich@gmail.com>
Date: Fri, 7 Feb 2025 23:54:21 -0800
Subject: [PATCH] Checkpoint UI tweaks

---
 .../chat/checkpoints/CheckpointMenu.tsx       | 110 +++++++++++-------
 .../chat/checkpoints/CheckpointSaved.tsx      |   2 +-
 webview-ui/src/components/ui/popover.tsx      |   8 +-
 3 files changed, 73 insertions(+), 47 deletions(-)

diff --git a/webview-ui/src/components/chat/checkpoints/CheckpointMenu.tsx b/webview-ui/src/components/chat/checkpoints/CheckpointMenu.tsx
index 657c292613..39141a6fbf 100644
--- a/webview-ui/src/components/chat/checkpoints/CheckpointMenu.tsx
+++ b/webview-ui/src/components/chat/checkpoints/CheckpointMenu.tsx
@@ -1,17 +1,9 @@
 import { useState, useEffect, useCallback } from "react"
-import { DotsHorizontalIcon } from "@radix-ui/react-icons"
-import { DropdownMenuItemProps } from "@radix-ui/react-dropdown-menu"
+import { CheckIcon, Cross2Icon } from "@radix-ui/react-icons"
 
 import { vscode } from "../../../utils/vscode"
 
-import {
-	Button,
-	DropdownMenu,
-	DropdownMenuTrigger,
-	DropdownMenuContent,
-	DropdownMenuItem,
-	DropdownMenuShortcut,
-} from "@/components/ui"
+import { Button, Popover, PopoverContent, PopoverTrigger } from "@/components/ui"
 
 type CheckpointMenuProps = {
 	ts: number
@@ -20,10 +12,8 @@ type CheckpointMenuProps = {
 
 export const CheckpointMenu = ({ ts, commitHash }: CheckpointMenuProps) => {
 	const [portalContainer, setPortalContainer] = useState<HTMLElement>()
-
-	const onTaskDiff = useCallback(() => {
-		vscode.postMessage({ type: "checkpointDiff", payload: { ts, commitHash, mode: "full" } })
-	}, [ts, commitHash])
+	const [isOpen, setIsOpen] = useState(false)
+	const [isConfirming, setIsConfirming] = useState(false)
 
 	const onCheckpointDiff = useCallback(() => {
 		vscode.postMessage({ type: "checkpointDiff", payload: { ts, commitHash, mode: "checkpoint" } })
@@ -31,10 +21,12 @@ export const CheckpointMenu = ({ ts, commitHash }: CheckpointMenuProps) => {
 
 	const onPreview = useCallback(() => {
 		vscode.postMessage({ type: "checkpointRestore", payload: { ts, commitHash, mode: "preview" } })
+		setIsOpen(false)
 	}, [ts, commitHash])
 
 	const onRestore = useCallback(() => {
 		vscode.postMessage({ type: "checkpointRestore", payload: { ts, commitHash, mode: "restore" } })
+		setIsOpen(false)
 	}, [ts, commitHash])
 
 	useEffect(() => {
@@ -47,34 +39,66 @@ export const CheckpointMenu = ({ ts, commitHash }: CheckpointMenuProps) => {
 	}, [])
 
 	return (
-		<DropdownMenu>
-			<DropdownMenuTrigger asChild>
-				<Button variant="ghost" size="icon">
-					<DotsHorizontalIcon />
-				</Button>
-			</DropdownMenuTrigger>
-			<DropdownMenuContent container={portalContainer} align="end">
-				<CheckpointMenuItem label="Checkpoint Diff" icon="diff-single" onClick={onCheckpointDiff} />
-				<CheckpointMenuItem label="Task Diff" icon="diff-multiple" onClick={onTaskDiff} />
-				<CheckpointMenuItem label="Preview" icon="open-preview" onClick={onPreview} />
-				<CheckpointMenuItem label="Restore" icon="history" onClick={onRestore} />
-			</DropdownMenuContent>
-		</DropdownMenu>
+		<div className="flex flex-row gap-1">
+			<Button variant="ghost" size="icon" onClick={onCheckpointDiff}>
+				<span className="codicon codicon-diff-single" />
+			</Button>
+			<Popover
+				open={isOpen}
+				onOpenChange={(open) => {
+					setIsOpen(open)
+					setIsConfirming(false)
+				}}>
+				<PopoverTrigger asChild>
+					<Button variant="ghost" size="icon">
+						<span className="codicon codicon-history" />
+					</Button>
+				</PopoverTrigger>
+				<PopoverContent align="end" container={portalContainer}>
+					<div className="flex flex-col gap-2">
+						<div className="flex flex-col gap-1 group hover:text-foreground">
+							<Button variant="secondary" onClick={onPreview}>
+								Restore Files
+							</Button>
+							<div className="text-muted transition-colors group-hover:text-foreground">
+								Restores your project's files back to a snapshot taken at this point.
+							</div>
+						</div>
+						<div className="flex flex-col gap-1 group hover:text-foreground">
+							<div className="flex flex-col gap-1 group hover:text-foreground">
+								{!isConfirming ? (
+									<Button variant="secondary" onClick={() => setIsConfirming(true)}>
+										Restore Files & Task
+									</Button>
+								) : (
+									<>
+										<Button variant="default" onClick={onRestore} className="grow">
+											<div className="flex flex-row gap-1">
+												<CheckIcon />
+												<div>Confirm</div>
+											</div>
+										</Button>
+										<Button variant="secondary" onClick={() => setIsConfirming(false)}>
+											<div className="flex flex-row gap-1">
+												<Cross2Icon />
+												<div>Cancel</div>
+											</div>
+										</Button>
+									</>
+								)}
+								{isConfirming ? (
+									<div className="text-destructive font-bold">This action cannot be undone.</div>
+								) : (
+									<div className="text-muted transition-colors group-hover:text-foreground">
+										Restores your project's files back to a snapshot taken at this point and deletes
+										all messages after this point.
+									</div>
+								)}
+							</div>
+						</div>
+					</div>
+				</PopoverContent>
+			</Popover>
+		</div>
 	)
 }
-
-type CheckpointMenuItemProps = DropdownMenuItemProps & {
-	label: React.ReactNode
-	icon: "diff-single" | "diff-multiple" | "open-preview" | "history"
-}
-
-const CheckpointMenuItem = ({ label, icon, ...props }: CheckpointMenuItemProps) => (
-	<DropdownMenuItem {...props}>
-		<div className="flex flex-row-reverse gap-1">
-			<div>{label}</div>
-			<DropdownMenuShortcut>
-				<span className={`codicon codicon-${icon}`} />
-			</DropdownMenuShortcut>
-		</div>
-	</DropdownMenuItem>
-)
diff --git a/webview-ui/src/components/chat/checkpoints/CheckpointSaved.tsx b/webview-ui/src/components/chat/checkpoints/CheckpointSaved.tsx
index 103f1d172c..5947b6a1fc 100644
--- a/webview-ui/src/components/chat/checkpoints/CheckpointSaved.tsx
+++ b/webview-ui/src/components/chat/checkpoints/CheckpointSaved.tsx
@@ -8,7 +8,7 @@ type CheckpointSavedProps = {
 export const CheckpointSaved = (props: CheckpointSavedProps) => (
 	<div className="flex items-center justify-between">
 		<div className="flex items-center gap-2">
-			<span className="codicon codicon-git-commit" />
+			<span className="codicon codicon-git-commit text-blue-400" />
 			<span className="font-bold">Checkpoint</span>
 		</div>
 		<CheckpointMenu {...props} />
diff --git a/webview-ui/src/components/ui/popover.tsx b/webview-ui/src/components/ui/popover.tsx
index 09c6ba659a..3ab3344689 100644
--- a/webview-ui/src/components/ui/popover.tsx
+++ b/webview-ui/src/components/ui/popover.tsx
@@ -11,9 +11,11 @@ const PopoverAnchor = PopoverPrimitive.Anchor
 
 const PopoverContent = React.forwardRef<
 	React.ElementRef<typeof PopoverPrimitive.Content>,
-	React.ComponentPropsWithoutRef<typeof PopoverPrimitive.Content>
->(({ className, align = "center", sideOffset = 4, ...props }, ref) => (
-	<PopoverPrimitive.Portal>
+	React.ComponentPropsWithoutRef<typeof PopoverPrimitive.Content> & {
+		container?: HTMLElement
+	}
+>(({ className, align = "center", sideOffset = 4, container, ...props }, ref) => (
+	<PopoverPrimitive.Portal container={container}>
 		<PopoverPrimitive.Content
 			ref={ref}
 			align={align}