Merge pull request #866 from RooVetGit/cte/checkpoints-ui-tweaks

Checkpoint UI tweaks

src/core/Cline.ts

@@ -3313,6 +3313,10 @@ export class Cline {
 			const commit = await service.saveCheckpoint(`Task: ${this.taskId}, Time: ${Date.now()}`)
 
 			if (commit?.commit) {
+				await this.providerRef
+					.deref()
+					?.postMessageToWebview({ type: "currentCheckpointUpdated", text: service.currentCheckpoint })
+
 				await this.say("checkpoint_saved", commit.commit)
 			}
 		} catch (err) {
@@ -3349,6 +3353,10 @@ export class Cline {
 			const service = await this.getCheckpointService()
 			await service.restoreCheckpoint(commitHash)
 
+			await this.providerRef
+				.deref()
+				?.postMessageToWebview({ type: "currentCheckpointUpdated", text: service.currentCheckpoint })
+
 			if (mode === "restore") {
 				await this.overwriteApiConversationHistory(
 					this.apiConversationHistory.filter((m) => !m.ts || m.ts < ts),

src/services/checkpoints/CheckpointService.ts

@@ -9,12 +9,6 @@ if (process.env.NODE_ENV !== "test") {
 	debug.enable("simple-git")
 }
 
-export interface Checkpoint {
-	hash: string
-	message: string
-	timestamp?: Date
-}
-
 export type CheckpointServiceOptions = {
 	taskId: string
 	git?: SimpleGit
@@ -60,6 +54,16 @@ export type CheckpointServiceOptions = {
  */
 
 export class CheckpointService {
+	private _currentCheckpoint?: string
+
+	public get currentCheckpoint() {
+		return this._currentCheckpoint
+	}
+
+	private set currentCheckpoint(value: string | undefined) {
+		this._currentCheckpoint = value
+	}
+
 	constructor(
 		public readonly taskId: string,
 		private readonly git: SimpleGit,
@@ -217,6 +221,8 @@ export class CheckpointService {
 				await this.popStash()
 			}
 
+			this.currentCheckpoint = commit.commit
+
 			return commit
 		} catch (err) {
 			this.log(`[saveCheckpoint] Failed to save checkpoint: ${err instanceof Error ? err.message : String(err)}`)
@@ -237,6 +243,7 @@ export class CheckpointService {
 		await this.ensureBranch(this.mainBranch)
 		await this.git.clean([CleanOptions.FORCE, CleanOptions.RECURSIVE])
 		await this.git.raw(["restore", "--source", commitHash, "--worktree", "--", "."])
+		this.currentCheckpoint = commitHash
 	}
 
 	public static async create({ taskId, git, baseDir, log = console.log }: CheckpointServiceOptions) {
@@ -291,7 +298,7 @@ export class CheckpointService {
 			// the checkpoint (i.e. the `git restore` command doesn't work
 			// for empty commits).
 			await fs.writeFile(path.join(baseDir, ".gitkeep"), "")
-			await git.add(".")
+			await git.add(".gitkeep")
 			const commit = await git.commit("Initial commit")
 
 			if (!commit.commit) {

src/services/checkpoints/__tests__/CheckpointService.test.ts

@@ -291,6 +291,7 @@ describe("CheckpointService", () => {
 			const baseDir = path.join(os.tmpdir(), `checkpoint-service-test2-${Date.now()}`)
 			await fs.mkdir(baseDir)
 			const newTestFile = path.join(baseDir, "test.txt")
+			await fs.writeFile(newTestFile, "Hello, world!")
 
 			const newGit = simpleGit(baseDir)
 			const initSpy = jest.spyOn(newGit, "init")
@@ -300,7 +301,6 @@ describe("CheckpointService", () => {
 			expect(initSpy).toHaveBeenCalled()
 
 			// Save a checkpoint: Hello, world!
-			await fs.writeFile(newTestFile, "Hello, world!")
 			const commit1 = await newService.saveCheckpoint("Hello, world!")
 			expect(commit1?.commit).toBeTruthy()
 			expect(await fs.readFile(newTestFile, "utf-8")).toBe("Hello, world!")

src/shared/ExtensionMessage.ts

@@ -42,6 +42,7 @@ export interface ExtensionMessage {
 		| "autoApprovalEnabled"
 		| "updateCustomMode"
 		| "deleteCustomMode"
+		| "currentCheckpointUpdated"
 	text?: string
 	action?:
 		| "chatButtonClicked"

webview-ui/src/components/chat/ChatRow.tsx

@@ -81,7 +81,7 @@ export const ChatRowContent = ({
 	isLast,
 	isStreaming,
 }: ChatRowContentProps) => {
-	const { mcpServers, alwaysAllowMcp } = useExtensionState()
+	const { mcpServers, alwaysAllowMcp, currentCheckpoint } = useExtensionState()
 	const [reasoningCollapsed, setReasoningCollapsed] = useState(false)
 
 	// Auto-collapse reasoning when new messages arrive
@@ -757,7 +757,13 @@ export const ChatRowContent = ({
 						</>
 					)
 				case "checkpoint_saved":
-					return <CheckpointSaved ts={message.ts!} commitHash={message.text!} />
+					return (
+						<CheckpointSaved
+							ts={message.ts!}
+							commitHash={message.text!}
+							currentCheckpointHash={currentCheckpoint}
+						/>
+					)
 				default:
 					return (
 						<>

webview-ui/src/components/chat/checkpoints/CheckpointMenu.tsx

@@ -1,29 +1,22 @@
 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
 	commitHash: string
+	currentCheckpointHash?: string
 }
 
-export const CheckpointMenu = ({ ts, commitHash }: CheckpointMenuProps) => {
+export const CheckpointMenu = ({ ts, commitHash, currentCheckpointHash }: CheckpointMenuProps) => {
 	const [portalContainer, setPortalContainer] = useState<HTMLElement>()
+	const [isOpen, setIsOpen] = useState(false)
+	const [isConfirming, setIsConfirming] = useState(false)
 
-	const onTaskDiff = useCallback(() => {
-		vscode.postMessage({ type: "checkpointDiff", payload: { ts, commitHash, mode: "full" } })
-	}, [ts, commitHash])
+	const isCurrent = currentCheckpointHash === commitHash
 
 	const onCheckpointDiff = useCallback(() => {
 		vscode.postMessage({ type: "checkpointDiff", payload: { ts, commitHash, mode: "checkpoint" } })
@@ -31,10 +24,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 +42,68 @@ 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">
+						{!isCurrent && (
+							<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>
-)

webview-ui/src/components/chat/checkpoints/CheckpointSaved.tsx

@@ -3,14 +3,20 @@ import { CheckpointMenu } from "./CheckpointMenu"
 type CheckpointSavedProps = {
 	ts: number
 	commitHash: string
+	currentCheckpointHash?: string
 }
 
-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="font-bold">Checkpoint</span>
+export const CheckpointSaved = (props: CheckpointSavedProps) => {
+	const isCurrent = props.currentCheckpointHash === props.commitHash
+
+	return (
+		<div className="flex items-center justify-between">
+			<div className="flex gap-2">
+				<span className="codicon codicon-git-commit text-blue-400" />
+				<span className="font-bold">Checkpoint</span>
+				{isCurrent && <span className="text-muted text-sm">Current</span>}
+			</div>
+			<CheckpointMenu {...props} />
 		</div>
-		<CheckpointMenu {...props} />
-	</div>
-)
+	)
+}

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}

webview-ui/src/context/ExtensionStateContext.tsx

@@ -26,6 +26,7 @@ export interface ExtensionStateContextType extends ExtensionState {
 	openRouterModels: Record<string, ModelInfo>
 	openAiModels: string[]
 	mcpServers: McpServer[]
+	currentCheckpoint?: string
 	filePaths: string[]
 	openedTabs: Array<{ label: string; isActive: boolean; path?: string }>
 	setApiConfiguration: (config: ApiConfiguration) => void
@@ -126,6 +127,7 @@ export const ExtensionStateContextProvider: React.FC<{ children: React.ReactNode
 
 	const [openAiModels, setOpenAiModels] = useState<string[]>([])
 	const [mcpServers, setMcpServers] = useState<McpServer[]>([])
+	const [currentCheckpoint, setCurrentCheckpoint] = useState<string>()
 
 	const setListApiConfigMeta = useCallback(
 		(value: ApiConfigMeta[]) => setState((prevState) => ({ ...prevState, listApiConfigMeta: value })),
@@ -241,6 +243,10 @@ export const ExtensionStateContextProvider: React.FC<{ children: React.ReactNode
 					setMcpServers(message.mcpServers ?? [])
 					break
 				}
+				case "currentCheckpointUpdated": {
+					setCurrentCheckpoint(message.text)
+					break
+				}
 				case "listApiConfig": {
 					setListApiConfigMeta(message.listApiConfig ?? [])
 					break
@@ -265,6 +271,7 @@ export const ExtensionStateContextProvider: React.FC<{ children: React.ReactNode
 		openRouterModels,
 		openAiModels,
 		mcpServers,
+		currentCheckpoint,
 		filePaths,
 		openedTabs,
 		soundVolume: state.soundVolume,