Changing checkpoint timing and ensuring checkpoints work (#6359)

* feat: Before requesting, ensure checkpoint is initialized

* Generate a checkpoint before modifying the code

* refactor: streamline checkpoint handling and enhance getCheckpoints method

* Blocked waiting for checkpoint initialization timing to change

* cancel checkpoint restore limit

* fix: ensure checkpoint service is undefined on initialization error and improve checkpoint diff handling

* refactor: simplify checkpoint service initialization and cleanup unused variables in CheckpointMenu

* fix: prevent race condition in checkpoint service initialization

- Only assign service to cline.checkpointService after successful initialization
- Add proper cleanup on initialization failure
- Prevents service from being in inconsistent state if Git check fails

* fix: remove checkpoint save from presentAssistantMessage for update_todo_list case

---------

Co-authored-by: Daniel Riccio <ricciodaniel98@gmail.com>

src/core/assistant-message/presentAssistantMessage.ts

@@ -25,7 +25,6 @@ import { switchModeTool } from "../tools/switchModeTool"
 import { attemptCompletionTool } from "../tools/attemptCompletionTool"
 import { newTaskTool } from "../tools/newTaskTool"
 
-import { checkpointSave } from "../checkpoints"
 import { updateTodoListTool } from "../tools/updateTodoListTool"
 
 import { formatResponse } from "../prompts/responses"
@@ -411,6 +410,7 @@ export async function presentAssistantMessage(cline: Task) {
 
 			switch (block.name) {
 				case "write_to_file":
+					await checkpointSaveAndMark(cline)
 					await writeToFileTool(cline, block, askApproval, handleError, pushToolResult, removeClosingTag)
 					break
 				case "update_todo_list":
@@ -430,8 +430,10 @@ export async function presentAssistantMessage(cline: Task) {
 					}
 
 					if (isMultiFileApplyDiffEnabled) {
+						await checkpointSaveAndMark(cline)
 						await applyDiffTool(cline, block, askApproval, handleError, pushToolResult, removeClosingTag)
 					} else {
+						await checkpointSaveAndMark(cline)
 						await applyDiffToolLegacy(
 							cline,
 							block,
@@ -444,9 +446,11 @@ export async function presentAssistantMessage(cline: Task) {
 					break
 				}
 				case "insert_content":
+					await checkpointSaveAndMark(cline)
 					await insertContentTool(cline, block, askApproval, handleError, pushToolResult, removeClosingTag)
 					break
 				case "search_and_replace":
+					await checkpointSaveAndMark(cline)
 					await searchAndReplaceTool(cline, block, askApproval, handleError, pushToolResult, removeClosingTag)
 					break
 				case "read_file":
@@ -527,14 +531,6 @@ export async function presentAssistantMessage(cline: Task) {
 			break
 	}
 
-	const recentlyModifiedFiles = cline.fileContextTracker.getAndClearCheckpointPossibleFile()
-
-	if (recentlyModifiedFiles.length > 0) {
-		// TODO: We can track what file changes were made and only
-		// checkpoint those files, this will be save storage.
-		await checkpointSave(cline)
-	}
-
 	// Seeing out of bounds is fine, it means that the next too call is being
 	// built up and ready to add to assistantMessageContent to present.
 	// When you see the UI inactive during this, it means that a tool is
@@ -583,3 +579,20 @@ export async function presentAssistantMessage(cline: Task) {
 		presentAssistantMessage(cline)
 	}
 }
+
+/**
+ * save checkpoint and mark done in the current streaming task.
+ * @param task The Task instance to checkpoint save and mark.
+ * @returns
+ */
+async function checkpointSaveAndMark(task: Task) {
+	if (task.currentStreamingDidCheckpoint) {
+		return
+	}
+	try {
+		await task.checkpointSave(true)
+		task.currentStreamingDidCheckpoint = true
+	} catch (error) {
+		console.error(`[Task#presentAssistantMessage] Error saving checkpoint: ${error.message}`, error)
+	}
+}

src/core/checkpoints/index.ts

@@ -16,18 +16,29 @@ import { DIFF_VIEW_URI_SCHEME } from "../../integrations/editor/DiffViewProvider
 
 import { CheckpointServiceOptions, RepoPerTaskCheckpointService } from "../../services/checkpoints"
 
-export function getCheckpointService(cline: Task) {
+export async function getCheckpointService(
+	cline: Task,
+	{ interval = 250, timeout = 15_000 }: { interval?: number; timeout?: number } = {},
+) {
 	if (!cline.enableCheckpoints) {
 		return undefined
 	}
 
 	if (cline.checkpointService) {
-		return cline.checkpointService
-	}
-
-	if (cline.checkpointServiceInitializing) {
-		console.log("[Task#getCheckpointService] checkpoint service is still initializing")
-		return undefined
+		if (cline.checkpointServiceInitializing) {
+			console.log("[Task#getCheckpointService] checkpoint service is still initializing")
+			const service = cline.checkpointService
+			await pWaitFor(
+				() => {
+					console.log("[Task#getCheckpointService] waiting for service to initialize")
+					return service.isInitialized
+				},
+				{ interval, timeout },
+			)
+			return service.isInitialized ? cline.checkpointService : undefined
+		} else {
+			return cline.checkpointService
+		}
 	}
 
 	const provider = cline.providerRef.deref()
@@ -69,15 +80,20 @@ export function getCheckpointService(cline: Task) {
 		}
 
 		const service = RepoPerTaskCheckpointService.create(options)
-
 		cline.checkpointServiceInitializing = true
 
 		// Check if Git is installed before initializing the service
-		// Note: This is intentionally fire-and-forget to match the original IIFE pattern
-		// The service is returned immediately while Git check happens asynchronously
-		checkGitInstallation(cline, service, log, provider)
-
-		return service
+		// Only assign the service after successful initialization
+		try {
+			await checkGitInstallation(cline, service, log, provider)
+			cline.checkpointService = service
+			return service
+		} catch (err) {
+			// Clean up on failure
+			cline.checkpointServiceInitializing = false
+			cline.enableCheckpoints = false
+			throw err
+		}
 	} catch (err) {
 		log(`[Task#getCheckpointService] ${err.message}`)
 		cline.enableCheckpoints = false
@@ -115,22 +131,7 @@ async function checkGitInstallation(
 		// Git is installed, proceed with initialization
 		service.on("initialize", () => {
 			log("[Task#getCheckpointService] service initialized")
-
-			try {
-				const isCheckpointNeeded =
-					typeof cline.clineMessages.find(({ say }) => say === "checkpoint_saved") === "undefined"
-
-				cline.checkpointService = service
-				cline.checkpointServiceInitializing = false
-
-				if (isCheckpointNeeded) {
-					log("[Task#getCheckpointService] no checkpoints found, saving initial checkpoint")
-					checkpointSave(cline)
-				}
-			} catch (err) {
-				log("[Task#getCheckpointService] caught error in on('initialize'), disabling checkpoints")
-				cline.enableCheckpoints = false
-			}
+			cline.checkpointServiceInitializing = false
 		})
 
 		service.on("checkpoint", ({ isFirst, fromHash: from, toHash: to }) => {
@@ -153,11 +154,12 @@ async function checkGitInstallation(
 		})
 
 		log("[Task#getCheckpointService] initializing shadow git")
-
-		service.initShadowGit().catch((err) => {
+		try {
+			await service.initShadowGit()
+		} catch (err) {
 			log(`[Task#getCheckpointService] initShadowGit -> ${err.message}`)
 			cline.enableCheckpoints = false
-		})
+		}
 	} catch (err) {
 		log(`[Task#getCheckpointService] Unexpected error during Git check: ${err.message}`)
 		console.error("Git check error:", err)
@@ -166,33 +168,8 @@ async function checkGitInstallation(
 	}
 }
 
-async function getInitializedCheckpointService(
-	cline: Task,
-	{ interval = 250, timeout = 15_000 }: { interval?: number; timeout?: number } = {},
-) {
-	const service = getCheckpointService(cline)
-
-	if (!service || service.isInitialized) {
-		return service
-	}
-
-	try {
-		await pWaitFor(
-			() => {
-				console.log("[Task#getCheckpointService] waiting for service to initialize")
-				return service.isInitialized
-			},
-			{ interval, timeout },
-		)
-
-		return service
-	} catch (err) {
-		return undefined
-	}
-}
-
 export async function checkpointSave(cline: Task, force = false) {
-	const service = getCheckpointService(cline)
+	const service = await getCheckpointService(cline)
 
 	if (!service) {
 		return
@@ -221,7 +198,7 @@ export type CheckpointRestoreOptions = {
 }
 
 export async function checkpointRestore(cline: Task, { ts, commitHash, mode }: CheckpointRestoreOptions) {
-	const service = await getInitializedCheckpointService(cline)
+	const service = await getCheckpointService(cline)
 
 	if (!service) {
 		return
@@ -289,7 +266,7 @@ export type CheckpointDiffOptions = {
 }
 
 export async function checkpointDiff(cline: Task, { ts, previousCommitHash, commitHash, mode }: CheckpointDiffOptions) {
-	const service = await getInitializedCheckpointService(cline)
+	const service = await getCheckpointService(cline)
 
 	if (!service) {
 		return
@@ -297,17 +274,19 @@ export async function checkpointDiff(cline: Task, { ts, previousCommitHash, comm
 
 	TelemetryService.instance.captureCheckpointDiffed(cline.taskId)
 
-	if (!previousCommitHash && mode === "checkpoint") {
-		const previousCheckpoint = cline.clineMessages
-			.filter(({ say }) => say === "checkpoint_saved")
-			.sort((a, b) => b.ts - a.ts)
-			.find((message) => message.ts < ts)
+	let prevHash = commitHash
+	let nextHash: string | undefined
 
-		previousCommitHash = previousCheckpoint?.text
+	const checkpoints = typeof service.getCheckpoints === "function" ? service.getCheckpoints() : []
+	const idx = checkpoints.indexOf(commitHash)
+	if (idx !== -1 && idx < checkpoints.length - 1) {
+		nextHash = checkpoints[idx + 1]
+	} else {
+		nextHash = undefined
 	}
 
 	try {
-		const changes = await service.getDiff({ from: previousCommitHash, to: commitHash })
+		const changes = await service.getDiff({ from: prevHash, to: nextHash })
 
 		if (!changes?.length) {
 			vscode.window.showInformationMessage("No changes found.")

src/core/task/Task.ts

@@ -240,6 +240,7 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 	isWaitingForFirstChunk = false
 	isStreaming = false
 	currentStreamingContentIndex = 0
+	currentStreamingDidCheckpoint = false
 	assistantMessageContent: AssistantMessageContent[] = []
 	presentAssistantMessageLocked = false
 	presentAssistantMessageHasPendingUpdates = false
@@ -1543,6 +1544,7 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 
 			// Reset streaming state.
 			this.currentStreamingContentIndex = 0
+			this.currentStreamingDidCheckpoint = false
 			this.assistantMessageContent = []
 			this.didCompleteReadingStream = false
 			this.userMessageContent = []

src/services/checkpoints/ShadowCheckpointService.ts

@@ -38,6 +38,10 @@ export abstract class ShadowCheckpointService extends EventEmitter {
 		return !!this.git
 	}
 
+	public getCheckpoints(): string[] {
+		return this._checkpoints.slice()
+	}
+
 	constructor(taskId: string, checkpointsDir: string, workspaceDir: string, log: (message: string) => void) {
 		super()
 

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

@@ -22,9 +22,6 @@ export const CheckpointMenu = ({ ts, commitHash, currentHash, checkpoint }: Chec
 	const portalContainer = useRooPortal("roo-portal")
 
 	const isCurrent = currentHash === commitHash
-	const isFirst = checkpoint.isFirst
-	const isDiffAvailable = !isFirst
-	const isRestoreAvailable = !isFirst || !isCurrent
 
 	const previousCommitHash = checkpoint?.from
 
@@ -47,78 +44,72 @@ export const CheckpointMenu = ({ ts, commitHash, currentHash, checkpoint }: Chec
 
 	return (
 		<div className="flex flex-row gap-1">
-			{isDiffAvailable && (
-				<StandardTooltip content={t("chat:checkpoint.menu.viewDiff")}>
-					<Button variant="ghost" size="icon" onClick={onCheckpointDiff}>
-						<span className="codicon codicon-diff-single" />
-					</Button>
+			<StandardTooltip content={t("chat:checkpoint.menu.viewDiff")}>
+				<Button variant="ghost" size="icon" onClick={onCheckpointDiff}>
+					<span className="codicon codicon-diff-single" />
+				</Button>
+			</StandardTooltip>
+			<Popover
+				open={isOpen}
+				onOpenChange={(open) => {
+					setIsOpen(open)
+					setIsConfirming(false)
+				}}>
+				<StandardTooltip content={t("chat:checkpoint.menu.restore")}>
+					<PopoverTrigger asChild>
+						<Button variant="ghost" size="icon">
+							<span className="codicon codicon-history" />
+						</Button>
+					</PopoverTrigger>
 				</StandardTooltip>
-			)}
-			{isRestoreAvailable && (
-				<Popover
-					open={isOpen}
-					onOpenChange={(open) => {
-						setIsOpen(open)
-						setIsConfirming(false)
-					}}>
-					<StandardTooltip content={t("chat:checkpoint.menu.restore")}>
-						<PopoverTrigger asChild>
-							<Button variant="ghost" size="icon">
-								<span className="codicon codicon-history" />
-							</Button>
-						</PopoverTrigger>
-					</StandardTooltip>
-					<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}>
-										{t("chat:checkpoint.menu.restoreFiles")}
+				<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}>
+									{t("chat:checkpoint.menu.restoreFiles")}
+								</Button>
+								<div className="text-muted transition-colors group-hover:text-foreground">
+									{t("chat:checkpoint.menu.restoreFilesDescription")}
+								</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)}>
+										{t("chat:checkpoint.menu.restoreFilesAndTask")}
 									</Button>
+								) : (
+									<>
+										<Button variant="default" onClick={onRestore} className="grow">
+											<div className="flex flex-row gap-1">
+												<CheckIcon />
+												<div>{t("chat:checkpoint.menu.confirm")}</div>
+											</div>
+										</Button>
+										<Button variant="secondary" onClick={() => setIsConfirming(false)}>
+											<div className="flex flex-row gap-1">
+												<Cross2Icon />
+												<div>{t("chat:checkpoint.menu.cancel")}</div>
+											</div>
+										</Button>
+									</>
+								)}
+								{isConfirming ? (
+									<div className="text-destructive font-bold">
+										{t("chat:checkpoint.menu.cannotUndo")}
+									</div>
+								) : (
 									<div className="text-muted transition-colors group-hover:text-foreground">
-										{t("chat:checkpoint.menu.restoreFilesDescription")}
+										{t("chat:checkpoint.menu.restoreFilesAndTaskDescription")}
 									</div>
-								</div>
-							)}
-							{!isFirst && (
-								<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)}>
-												{t("chat:checkpoint.menu.restoreFilesAndTask")}
-											</Button>
-										) : (
-											<>
-												<Button variant="default" onClick={onRestore} className="grow">
-													<div className="flex flex-row gap-1">
-														<CheckIcon />
-														<div>{t("chat:checkpoint.menu.confirm")}</div>
-													</div>
-												</Button>
-												<Button variant="secondary" onClick={() => setIsConfirming(false)}>
-													<div className="flex flex-row gap-1">
-														<Cross2Icon />
-														<div>{t("chat:checkpoint.menu.cancel")}</div>
-													</div>
-												</Button>
-											</>
-										)}
-										{isConfirming ? (
-											<div className="text-destructive font-bold">
-												{t("chat:checkpoint.menu.cannotUndo")}
-											</div>
-										) : (
-											<div className="text-muted transition-colors group-hover:text-foreground">
-												{t("chat:checkpoint.menu.restoreFilesAndTaskDescription")}
-											</div>
-										)}
-									</div>
-								</div>
-							)}
+								)}
+							</div>
 						</div>
-					</PopoverContent>
-				</Popover>
-			)}
+					</div>
+				</PopoverContent>
+			</Popover>
 		</div>
 	)
 }