Merge pull request #945 from RooVetGit/cte/checkpoints-logging

Checkpoints logging tweaks + defensive catches

src/core/Cline.ts

@@ -370,7 +370,13 @@ export class Cline {
 		this.askResponseImages = images
 	}
 
-	async say(type: ClineSay, text?: string, images?: string[], partial?: boolean): Promise<undefined> {
+	async say(
+		type: ClineSay,
+		text?: string,
+		images?: string[],
+		partial?: boolean,
+		checkpoint?: Record<string, unknown>,
+	): Promise<undefined> {
 		if (this.abort) {
 			throw new Error("Roo Code instance aborted")
 		}
@@ -423,7 +429,7 @@ export class Cline {
 			// this is a new non-partial message, so add it like normal
 			const sayTs = Date.now()
 			this.lastMessageTs = sayTs
-			await this.addToClineMessages({ ts: sayTs, type: "say", say: type, text, images })
+			await this.addToClineMessages({ ts: sayTs, type: "say", say: type, text, images, checkpoint })
 			await this.providerRef.deref()?.postStateToWebview()
 		}
 	}
@@ -2747,6 +2753,13 @@ export class Cline {
 		// get previous api req's index to check token usage and determine if we need to truncate conversation history
 		const previousApiReqIndex = findLastIndex(this.clineMessages, (m) => m.say === "api_req_started")
 
+		// Save checkpoint if this is the first API request.
+		const isFirstRequest = this.clineMessages.filter((m) => m.say === "api_req_started").length === 0
+
+		if (isFirstRequest) {
+			await this.checkpointSave()
+		}
+
 		// getting verbose details is an expensive operation, it uses globby to top-down build file structure of project which for large projects can take a few seconds
 		// for the best UX we show a placeholder api_req_started message with a loading spinner as this happens
 		await this.say(
@@ -3288,12 +3301,7 @@ export class Cline {
 				]),
 			)
 		} catch (err) {
-			this.providerRef
-				.deref()
-				?.log(
-					`[checkpointDiff] Encountered unexpected error: $${err instanceof Error ? err.message : String(err)}`,
-				)
-
+			this.providerRef.deref()?.log("[checkpointDiff] disabling checkpoints for this task")
 			this.checkpointsEnabled = false
 		}
 	}
@@ -3304,6 +3312,7 @@ export class Cline {
 		}
 
 		try {
+			const isFirst = !this.checkpointService
 			const service = await this.getCheckpointService()
 			const commit = await service.saveCheckpoint(`Task: ${this.taskId}, Time: ${Date.now()}`)
 
@@ -3312,15 +3321,17 @@ export class Cline {
 					.deref()
 					?.postMessageToWebview({ type: "currentCheckpointUpdated", text: service.currentCheckpoint })
 
-				await this.say("checkpoint_saved", commit.commit)
+				// Checkpoint metadata required by the UI.
+				const checkpoint = {
+					isFirst,
+					from: service.baseCommitHash,
+					to: commit.commit,
+				}
+
+				await this.say("checkpoint_saved", commit.commit, undefined, undefined, checkpoint)
 			}
 		} catch (err) {
-			this.providerRef
-				.deref()
-				?.log(
-					`[checkpointSave] Encountered unexpected error: $${err instanceof Error ? err.message : String(err)}`,
-				)
-
+			this.providerRef.deref()?.log("[checkpointSave] disabling checkpoints for this task")
 			this.checkpointsEnabled = false
 		}
 	}
@@ -3390,12 +3401,7 @@ export class Cline {
 			// Cline instance.
 			this.providerRef.deref()?.cancelTask()
 		} catch (err) {
-			this.providerRef
-				.deref()
-				?.log(
-					`[restoreCheckpoint] Encountered unexpected error: $${err instanceof Error ? err.message : String(err)}`,
-				)
-
+			this.providerRef.deref()?.log("[checkpointRestore] disabling checkpoints for this task")
 			this.checkpointsEnabled = false
 		}
 	}

src/services/checkpoints/CheckpointService.ts

@@ -132,22 +132,64 @@ export class CheckpointService {
 		stashSha: string
 		force?: boolean
 	}) {
-		if (force) {
-			await this.git.checkout(["-f", this.mainBranch])
-		} else {
-			await this.git.checkout(this.mainBranch)
+		let currentBranch = await this.git.revparse(["--abbrev-ref", "HEAD"])
+
+		if (currentBranch !== this.mainBranch) {
+			if (force) {
+				try {
+					await this.git.checkout(["-f", this.mainBranch])
+				} catch (err) {
+					this.log(
+						`[restoreMain] failed to force checkout ${this.mainBranch}: ${err instanceof Error ? err.message : String(err)}`,
+					)
+				}
+			} else {
+				try {
+					await this.git.checkout(this.mainBranch)
+				} catch (err) {
+					this.log(
+						`[restoreMain] failed to checkout ${this.mainBranch}: ${err instanceof Error ? err.message : String(err)}`,
+					)
+
+					// Escalate to a forced checkout if we can't checkout the
+					// main branch under normal circumstances.
+					currentBranch = await this.git.revparse(["--abbrev-ref", "HEAD"])
+
+					if (currentBranch !== this.mainBranch) {
+						await this.git.checkout(["-f", this.mainBranch]).catch(() => {})
+					}
+				}
+			}
+		}
+
+		currentBranch = await this.git.revparse(["--abbrev-ref", "HEAD"])
+
+		if (currentBranch !== this.mainBranch) {
+			throw new Error(`Unable to restore ${this.mainBranch}`)
 		}
 
 		if (stashSha) {
 			this.log(`[restoreMain] applying stash ${stashSha}`)
-			await this.git.raw(["stash", "apply", "--index", stashSha])
+
+			try {
+				await this.git.raw(["stash", "apply", "--index", stashSha])
+			} catch (err) {
+				this.log(`[restoreMain] Failed to apply stash: ${err instanceof Error ? err.message : String(err)}`)
+			}
 		}
 
-		this.log(`[restoreMain] restoring from ${branch}`)
-		await this.git.raw(["restore", "--source", branch, "--worktree", "--", "."])
+		this.log(`[restoreMain] restoring from ${branch} branch`)
+
+		try {
+			await this.git.raw(["restore", "--source", branch, "--worktree", "--", "."])
+		} catch (err) {
+			this.log(`[restoreMain] Failed to restore branch: ${err instanceof Error ? err.message : String(err)}`)
+		}
 	}
 
 	public async saveCheckpoint(message: string) {
+		const startTime = Date.now()
+
 		await this.ensureBranch(this.mainBranch)
 
 		const stashSha = (await this.git.raw(["stash", "create"])).trim()
@@ -172,15 +214,13 @@ export class CheckpointService {
 		 */
 		try {
 			await this.git.add(["-A"])
-			const status = await this.git.status()
-			this.log(`[saveCheckpoint] status: ${JSON.stringify(status)}`)
 		} catch (err) {
-			await this.git.checkout(["-f", this.mainBranch])
-			await this.git.branch(["-D", stashBranch]).catch(() => {})
-
-			throw new Error(
-				`[saveCheckpoint] Failed in stage stash phase: ${err instanceof Error ? err.message : String(err)}`,
+			this.log(
+				`[saveCheckpoint] failed in stage stash phase: ${err instanceof Error ? err.message : String(err)}`,
 			)
+			await this.restoreMain({ branch: stashBranch, stashSha, force: true })
+			await this.git.branch(["-D", stashBranch]).catch(() => {})
+			throw err
 		}
 
 		/**
@@ -192,17 +232,25 @@ export class CheckpointService {
 		 *   - UNDO: Create branch
 		 *   - UNDO: Change branch
 		 */
+		let stashCommit
+
 		try {
-			// TODO: Add a test to see if empty commits break this.
-			const tempCommit = await this.git.commit(message, undefined, { "--no-verify": null })
-			this.log(`[saveCheckpoint] tempCommit: ${message} -> ${JSON.stringify(tempCommit)}`)
+			stashCommit = await this.git.commit(message, undefined, { "--no-verify": null })
+			this.log(`[saveCheckpoint] stashCommit: ${message} -> ${JSON.stringify(stashCommit)}`)
 		} catch (err) {
-			await this.git.checkout(["-f", this.mainBranch])
+			this.log(
+				`[saveCheckpoint] failed in stash commit phase: ${err instanceof Error ? err.message : String(err)}`,
+			)
+			await this.restoreMain({ branch: stashBranch, stashSha, force: true })
 			await this.git.branch(["-D", stashBranch]).catch(() => {})
+			throw err
+		}
 
-			throw new Error(
-				`[saveCheckpoint] Failed in stash commit phase: ${err instanceof Error ? err.message : String(err)}`,
-			)
+		if (!stashCommit) {
+			this.log("[saveCheckpoint] no stash commit")
+			await this.restoreMain({ branch: stashBranch, stashSha })
+			await this.git.branch(["-D", stashBranch])
+			return undefined
 		}
 
 		/**
@@ -219,12 +267,10 @@ export class CheckpointService {
 		try {
 			diff = await this.git.diff([latestSha, stashBranch])
 		} catch (err) {
+			this.log(`[saveCheckpoint] failed in diff phase: ${err instanceof Error ? err.message : String(err)}`)
 			await this.restoreMain({ branch: stashBranch, stashSha, force: true })
 			await this.git.branch(["-D", stashBranch]).catch(() => {})
-
-			throw new Error(
-				`[saveCheckpoint] Failed in diff phase: ${err instanceof Error ? err.message : String(err)}`,
-			)
+			throw err
 		}
 
 		if (!diff) {
@@ -249,12 +295,10 @@ export class CheckpointService {
 			await this.git.reset(["--hard", this.mainBranch])
 			this.log(`[saveCheckpoint] reset ${this.hiddenBranch}`)
 		} catch (err) {
+			this.log(`[saveCheckpoint] failed in reset phase: ${err instanceof Error ? err.message : String(err)}`)
 			await this.restoreMain({ branch: stashBranch, stashSha, force: true })
 			await this.git.branch(["-D", stashBranch]).catch(() => {})
-
-			throw new Error(
-				`[saveCheckpoint] Failed in reset phase: ${err instanceof Error ? err.message : String(err)}`,
-			)
+			throw err
 		}
 
 		/**
@@ -289,25 +333,33 @@ export class CheckpointService {
 			this.currentCheckpoint = commit
 			this.log(`[saveCheckpoint] cherry-pick commit = ${commit}`)
 		} catch (err) {
+			this.log(
+				`[saveCheckpoint] failed in cherry pick phase: ${err instanceof Error ? err.message : String(err)}`,
+			)
 			await this.git.reset(["--hard", latestSha]).catch(() => {})
 			await this.restoreMain({ branch: stashBranch, stashSha, force: true })
 			await this.git.branch(["-D", stashBranch]).catch(() => {})
-
-			throw new Error(
-				`[saveCheckpoint] Failed in cherry pick phase: ${err instanceof Error ? err.message : String(err)}`,
-			)
+			throw err
 		}
 
 		await this.restoreMain({ branch: stashBranch, stashSha })
 		await this.git.branch(["-D", stashBranch])
 
+		// We've gotten reports that checkpoints can be slow in some cases, so
+		// we'll log the duration of the checkpoint save.
+		const duration = Date.now() - startTime
+		this.log(`[saveCheckpoint] saved checkpoint ${commit} in ${duration}ms`)
+
 		return { commit }
 	}
 
 	public async restoreCheckpoint(commitHash: string) {
+		const startTime = Date.now()
 		await this.ensureBranch(this.mainBranch)
 		await this.git.clean([CleanOptions.FORCE, CleanOptions.RECURSIVE])
 		await this.git.raw(["restore", "--source", commitHash, "--worktree", "--", "."])
+		const duration = Date.now() - startTime
+		this.log(`[restoreCheckpoint] restored checkpoint ${commitHash} in ${duration}ms`)
 		this.currentCheckpoint = commitHash
 	}
 

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

@@ -210,9 +210,12 @@ describe("CheckpointService", () => {
 		})
 
 		it("does not create a checkpoint if there are no pending changes", async () => {
+			const commit0 = await service.saveCheckpoint("Zeroth checkpoint")
+			expect(commit0?.commit).toBeFalsy()
+
 			await fs.writeFile(testFile, "Ahoy, world!")
-			const commit = await service.saveCheckpoint("First checkpoint")
-			expect(commit?.commit).toBeTruthy()
+			const commit1 = await service.saveCheckpoint("First checkpoint")
+			expect(commit1?.commit).toBeTruthy()
 
 			const commit2 = await service.saveCheckpoint("Second checkpoint")
 			expect(commit2?.commit).toBeFalsy()

src/shared/ExtensionMessage.ts

@@ -139,6 +139,7 @@ export interface ClineMessage {
 	partial?: boolean
 	reasoning?: string
 	conversationHistoryIndex?: number
+	checkpoint?: Record<string, unknown>
 }
 
 export type ClineAsk =

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

@@ -761,6 +761,7 @@ export const ChatRowContent = ({
 						<CheckpointSaved
 							ts={message.ts!}
 							commitHash={message.text!}
+							checkpoint={message.checkpoint}
 							currentCheckpointHash={currentCheckpoint}
 						/>
 					)

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

@@ -1,17 +1,19 @@
 import { useState, useEffect, useCallback } from "react"
 import { CheckIcon, Cross2Icon } from "@radix-ui/react-icons"
 
-import { vscode } from "../../../utils/vscode"
-
 import { Button, Popover, PopoverContent, PopoverTrigger } from "@/components/ui"
 
+import { vscode } from "../../../utils/vscode"
+import { Checkpoint } from "./schema"
+
 type CheckpointMenuProps = {
 	ts: number
 	commitHash: string
+	checkpoint?: Checkpoint
 	currentCheckpointHash?: string
 }
 
-export const CheckpointMenu = ({ ts, commitHash, currentCheckpointHash }: CheckpointMenuProps) => {
+export const CheckpointMenu = ({ ts, commitHash, checkpoint, currentCheckpointHash }: CheckpointMenuProps) => {
 	const [portalContainer, setPortalContainer] = useState<HTMLElement>()
 	const [isOpen, setIsOpen] = useState(false)
 	const [isConfirming, setIsConfirming] = useState(false)
@@ -43,9 +45,11 @@ export const CheckpointMenu = ({ ts, commitHash, currentCheckpointHash }: Checkp
 
 	return (
 		<div className="flex flex-row gap-1">
-			<Button variant="ghost" size="icon" onClick={onCheckpointDiff} title="View Diff">
-				<span className="codicon codicon-diff-single" />
-			</Button>
+			{!checkpoint?.isFirst && (
+				<Button variant="ghost" size="icon" onClick={onCheckpointDiff} title="View Diff">
+					<span className="codicon codicon-diff-single" />
+				</Button>
+			)}
 			<Popover
 				open={isOpen}
 				onOpenChange={(open) => {

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

@@ -1,22 +1,37 @@
+import { useMemo } from "react"
+
 import { CheckpointMenu } from "./CheckpointMenu"
+import { checkpointSchema } from "./schema"
 
 type CheckpointSavedProps = {
 	ts: number
 	commitHash: string
+	checkpoint?: Record<string, unknown>
 	currentCheckpointHash?: string
 }
 
-export const CheckpointSaved = (props: CheckpointSavedProps) => {
+export const CheckpointSaved = ({ checkpoint, ...props }: CheckpointSavedProps) => {
 	const isCurrent = props.currentCheckpointHash === props.commitHash
 
+	const metadata = useMemo(() => {
+		if (!checkpoint) {
+			return undefined
+		}
+
+		const result = checkpointSchema.safeParse(checkpoint)
+		return result.success ? result.data : undefined
+	}, [checkpoint])
+
+	const isFirst = !!metadata?.isFirst
+
 	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>
+				<span className="font-bold">{isFirst ? "Initial Checkpoint" : "Checkpoint"}</span>
 				{isCurrent && <span className="text-muted text-sm">Current</span>}
 			</div>
-			<CheckpointMenu {...props} />
+			<CheckpointMenu {...props} checkpoint={metadata} />
 		</div>
 	)
 }

webview-ui/src/components/chat/checkpoints/schema.ts

@@ -0,0 +1,9 @@
+import { z } from "zod"
+
+export const checkpointSchema = z.object({
+	isFirst: z.boolean(),
+	from: z.string(),
+	to: z.string(),
+})
+
+export type Checkpoint = z.infer<typeof checkpointSchema>