Add a repo-per-workspace checkpoints services in addition to repo-per-task

Author: cte

src/core/Cline.ts

@@ -1,6 +1,6 @@
 import { Anthropic } from "@anthropic-ai/sdk"
 import cloneDeep from "clone-deep"
-import { DiffStrategy, getDiffStrategy, UnifiedDiffStrategy } from "./diff/DiffStrategy"
+import { DiffStrategy, getDiffStrategy } from "./diff/DiffStrategy"
 import { validateToolUse, isToolAllowedForMode, ToolName } from "./mode-validator"
 import delay from "delay"
 import fs from "fs/promises"
@@ -13,7 +13,11 @@ import * as vscode from "vscode"
 import { ApiHandler, buildApiHandler } from "../api"
 import { ApiStream } from "../api/transform/stream"
 import { DIFF_VIEW_URI_SCHEME, DiffViewProvider } from "../integrations/editor/DiffViewProvider"
-import { ShadowCheckpointService } from "../services/checkpoints/ShadowCheckpointService"
+import {
+	CheckpointServiceOptions,
+	RepoPerTaskCheckpointService,
+	RepoPerWorkspaceCheckpointService,
+} from "../services/checkpoints"
 import { findToolName, formatContentBlockToMarkdown } from "../integrations/misc/export-markdown"
 import {
 	extractTextFromFile,
@@ -77,6 +81,7 @@ export type ClineOptions = {
 	customInstructions?: string
 	enableDiff?: boolean
 	enableCheckpoints?: boolean
+	checkpointStorage?: "task" | "workspace"
 	fuzzyMatchThreshold?: number
 	task?: string
 	images?: string[]
@@ -115,8 +120,9 @@ export class Cline {
 	isInitialized = false
 
 	// checkpoints
-	enableCheckpoints: boolean = false
-	private checkpointService?: ShadowCheckpointService
+	private enableCheckpoints: boolean
+	private checkpointStorage: "task" | "workspace"
+	private checkpointService?: RepoPerTaskCheckpointService | RepoPerWorkspaceCheckpointService
 
 	// streaming
 	isWaitingForFirstChunk = false
@@ -136,7 +142,8 @@ export class Cline {
 		apiConfiguration,
 		customInstructions,
 		enableDiff,
-		enableCheckpoints,
+		enableCheckpoints = false,
+		checkpointStorage = "workspace",
 		fuzzyMatchThreshold,
 		task,
 		images,
@@ -160,7 +167,8 @@ export class Cline {
 		this.fuzzyMatchThreshold = fuzzyMatchThreshold ?? 1.0
 		this.providerRef = new WeakRef(provider)
 		this.diffViewProvider = new DiffViewProvider(cwd)
-		this.enableCheckpoints = enableCheckpoints ?? false
+		this.enableCheckpoints = enableCheckpoints
+		this.checkpointStorage = checkpointStorage
 
 		// Initialize diffStrategy based on current state
 		this.updateDiffStrategy(Experiments.isEnabled(experiments ?? {}, EXPERIMENT_IDS.DIFF_STRATEGY))
@@ -747,7 +755,8 @@ export class Cline {
 	}
 
 	private async initiateTaskLoop(userContent: UserContent): Promise<void> {
-		this.initializeCheckpoints()
+		// Kicks off the checkpoints initialization process in the background.
+		this.getCheckpointService()
 
 		let nextUserContent = userContent
 		let includeFileDetails = true
@@ -3352,9 +3361,13 @@ export class Cline {
 
 	// Checkpoints
 
-	private async initializeCheckpoints() {
+	private getCheckpointService() {
 		if (!this.enableCheckpoints) {
-			return
+			return undefined
+		}
+
+		if (this.checkpointService) {
+			return this.checkpointService
 		}
 
 		const log = (message: string) => {
@@ -3368,45 +3381,36 @@ export class Cline {
 		}
 
 		try {
-			if (this.checkpointService) {
-				log("[Cline#initializeCheckpoints] checkpointService already initialized")
-				return
-			}
-
 			const workspaceDir = vscode.workspace.workspaceFolders?.map((folder) => folder.uri.fsPath).at(0)
 
 			if (!workspaceDir) {
 				log("[Cline#initializeCheckpoints] workspace folder not found, disabling checkpoints")
 				this.enableCheckpoints = false
-				return
+				return undefined
 			}
 
-			const shadowDir = this.providerRef.deref()?.context.globalStorageUri.fsPath
+			const globalStorageDir = this.providerRef.deref()?.context.globalStorageUri.fsPath
 
-			if (!shadowDir) {
-				log("[Cline#initializeCheckpoints] shadowDir not found, disabling checkpoints")
+			if (!globalStorageDir) {
+				log("[Cline#initializeCheckpoints] globalStorageDir not found, disabling checkpoints")
 				this.enableCheckpoints = false
-				return
+				return undefined
 			}
 
-			const service = await ShadowCheckpointService.create({ taskId: this.taskId, workspaceDir, shadowDir, log })
-
-			if (!service) {
-				log("[Cline#initializeCheckpoints] failed to create checkpoint service, disabling checkpoints")
-				this.enableCheckpoints = false
-				return
+			const options: CheckpointServiceOptions = {
+				taskId: this.taskId,
+				workspaceDir,
+				shadowDir: globalStorageDir,
+				log,
 			}
 
-			service.on("initialize", ({ workspaceDir, created, duration }) => {
-				try {
-					if (created) {
-						log(`[Cline#initializeCheckpoints] created new shadow repo (${workspaceDir}) in ${duration}ms`)
-					} else {
-						log(
-							`[Cline#initializeCheckpoints] found existing shadow repo (${workspaceDir}) in ${duration}ms`,
-						)
-					}
+			const service =
+				this.checkpointStorage === "task"
+					? RepoPerTaskCheckpointService.create(options)
+					: RepoPerWorkspaceCheckpointService.create(options)
 
+			service.on("initialize", () => {
+				try {
 					this.checkpointService = service
 					this.checkpointSave()
 				} catch (err) {
@@ -3417,7 +3421,6 @@ export class Cline {
 
 			service.on("checkpoint", ({ isFirst, fromHash: from, toHash: to }) => {
 				try {
-					log(`[Cline#initializeCheckpoints] ${isFirst ? "initial" : "incremental"} checkpoint saved: ${to}`)
 					this.providerRef.deref()?.postMessageToWebview({ type: "currentCheckpointUpdated", text: to })
 
 					this.say("checkpoint_saved", to, undefined, undefined, { isFirst, from, to }).catch((e) =>
@@ -3430,9 +3433,35 @@ export class Cline {
 			})
 
 			service.initShadowGit()
+			return service
 		} catch (err) {
-			log("[Cline#initializeCheckpoints] caught error in initializeCheckpoints(), disabling checkpoints")
+			log("[Cline#initializeCheckpoints] caught unexpected error, disabling checkpoints")
 			this.enableCheckpoints = false
+			return undefined
+		}
+	}
+
+	private async getInitializedCheckpointService({
+		interval = 250,
+		timeout = 15_000,
+	}: { interval?: number; timeout?: number } = {}) {
+		const service = this.getCheckpointService()
+
+		if (!service || service.isInitialized) {
+			return service
+		}
+
+		try {
+			await pWaitFor(
+				() => {
+					console.log("[Cline#getCheckpointService] waiting for service to initialize")
+					return service.isInitialized
+				},
+				{ interval, timeout },
+			)
+			return service
+		} catch (err) {
+			return undefined
 		}
 	}
 
@@ -3445,7 +3474,9 @@ export class Cline {
 		commitHash: string
 		mode: "full" | "checkpoint"
 	}) {
-		if (!this.checkpointService || !this.enableCheckpoints) {
+		const service = await this.getInitializedCheckpointService()
+
+		if (!service) {
 			return
 		}
 
@@ -3461,7 +3492,7 @@ export class Cline {
 		}
 
 		try {
-			const changes = await this.checkpointService.getDiff({ from: previousCommitHash, to: commitHash })
+			const changes = await service.getDiff({ from: previousCommitHash, to: commitHash })
 
 			if (!changes?.length) {
 				vscode.window.showInformationMessage("No changes found.")
@@ -3488,12 +3519,22 @@ export class Cline {
 	}
 
 	public checkpointSave() {
-		if (!this.checkpointService || !this.enableCheckpoints) {
+		const service = this.getCheckpointService()
+
+		if (!service) {
+			return
+		}
+
+		if (!service.isInitialized) {
+			this.providerRef
+				.deref()
+				?.log("[checkpointSave] checkpoints didn't initialize in time, disabling checkpoints for this task")
+			this.enableCheckpoints = false
 			return
 		}
 
 		// Start the checkpoint process in the background.
-		this.checkpointService.saveCheckpoint(`Task: ${this.taskId}, Time: ${Date.now()}`)
+		service.saveCheckpoint(`Task: ${this.taskId}, Time: ${Date.now()}`)
 	}
 
 	public async checkpointRestore({
@@ -3505,7 +3546,9 @@ export class Cline {
 		commitHash: string
 		mode: "preview" | "restore"
 	}) {
-		if (!this.checkpointService || !this.enableCheckpoints) {
+		const service = await this.getInitializedCheckpointService()
+
+		if (!service) {
 			return
 		}
 
@@ -3516,7 +3559,7 @@ export class Cline {
 		}
 
 		try {
-			await this.checkpointService.restoreCheckpoint(commitHash)
+			await service.restoreCheckpoint(commitHash)
 
 			await this.providerRef.deref()?.postMessageToWebview({ type: "currentCheckpointUpdated", text: commitHash })
 

src/core/webview/ClineProvider.ts

@@ -887,6 +887,8 @@ export class ClineProvider implements vscode.WebviewViewProvider {
 
 						if (result.success) {
 							await this.cline?.checkpointDiff(result.data)
+						} else {
+							vscode.window.showErrorMessage("Invalid checkpoint diff payload.")
 						}
 
 						break
@@ -907,6 +909,8 @@ export class ClineProvider implements vscode.WebviewViewProvider {
 							} catch (error) {
 								vscode.window.showErrorMessage("Failed to restore checkpoint.")
 							}
+						} else {
+							vscode.window.showErrorMessage("Invalid checkpoint restore payload.")
 						}
 
 						break

src/services/checkpoints/RepoPerTaskCheckpointService.ts

@@ -0,0 +1,15 @@
+import * as path from "path"
+
+import { CheckpointServiceOptions } from "./types"
+import { ShadowCheckpointService } from "./ShadowCheckpointService"
+
+export class RepoPerTaskCheckpointService extends ShadowCheckpointService {
+	public static create({ taskId, workspaceDir, shadowDir, log = console.log }: CheckpointServiceOptions) {
+		return new RepoPerTaskCheckpointService(
+			taskId,
+			path.join(shadowDir, "tasks", taskId, "checkpoints"),
+			workspaceDir,
+			log,
+		)
+	}
+}

src/services/checkpoints/RepoPerWorkspaceCheckpointService.ts

@@ -0,0 +1,53 @@
+import * as path from "path"
+import crypto from "crypto"
+
+import { CheckpointServiceOptions } from "./types"
+import { ShadowCheckpointService } from "./ShadowCheckpointService"
+
+export class RepoPerWorkspaceCheckpointService extends ShadowCheckpointService {
+	private async checkoutTaskBranch() {
+		if (!this.git) {
+			throw new Error("Shadow git repo not initialized")
+		}
+
+		const startTime = Date.now()
+
+		const branch = `roo-${this.taskId}`
+		this.log(`[${this.constructor.name}#checkoutTaskBranch] searching branches for ${branch}`)
+		const branches = await this.git.branchLocal()
+
+		if (!branches.all.includes(branch)) {
+			this.log(`[${this.constructor.name}#checkoutTaskBranch] creating ${branch}`)
+			await this.git.checkoutLocalBranch(branch)
+		} else {
+			this.log(`[${this.constructor.name}#checkoutTaskBranch] checking out ${branch}`)
+			await this.git.checkout(branch)
+		}
+
+		const result = await this.git.revparse(["--abbrev-ref", "HEAD"])
+
+		const duration = Date.now() - startTime
+		this.log(`[${this.constructor.name}#checkoutTaskBranch] checked out ${result} in ${duration}ms`)
+	}
+
+	override async initShadowGit() {
+		await super.initShadowGit()
+		await this.checkoutTaskBranch()
+	}
+
+	override async restoreCheckpoint(commitHash: string) {
+		await this.checkoutTaskBranch()
+		await super.restoreCheckpoint(commitHash)
+	}
+
+	public static create({ taskId, workspaceDir, shadowDir, log = console.log }: CheckpointServiceOptions) {
+		const workspaceHash = crypto.createHash("sha256").update(workspaceDir).digest("hex").toString()
+
+		return new RepoPerWorkspaceCheckpointService(
+			taskId,
+			path.join(shadowDir, "checkpoints", workspaceHash),
+			workspaceDir,
+			log,
+		)
+	}
+}