Merge pull request #958 from RooVetGit/cte/shadow-repo-checkpoints

Shadow repo checkpoints

Author: cte

package-lock.json

@@ -314,6 +314,7 @@
 		"diff-match-patch": "^1.0.5",
 		"fast-deep-equal": "^3.1.3",
 		"fastest-levenshtein": "^1.0.16",
+		"get-folder-size": "^5.0.0",
 		"globby": "^14.0.2",
 		"isbinaryfile": "^5.0.2",
 		"mammoth": "^1.8.0",
@@ -322,6 +323,7 @@
 		"os-name": "^6.0.0",
 		"p-wait-for": "^5.0.2",
 		"pdf-parse": "^1.1.1",
+		"pretty-bytes": "^6.1.1",
 		"puppeteer-chromium-resolver": "^23.0.0",
 		"puppeteer-core": "^23.4.0",
 		"serialize-error": "^11.0.3",
@@ -352,17 +354,17 @@
 		"@vscode/test-cli": "^0.0.9",
 		"@vscode/test-electron": "^2.4.0",
 		"esbuild": "^0.24.0",
-		"mkdirp": "^3.0.1",
-		"rimraf": "^6.0.1",
 		"eslint": "^8.57.0",
 		"glob": "^11.0.1",
 		"husky": "^9.1.7",
 		"jest": "^29.7.0",
 		"jest-simple-dot-reporter": "^1.0.5",
 		"lint-staged": "^15.2.11",
+		"mkdirp": "^3.0.1",
 		"mocha": "^11.1.0",
 		"npm-run-all": "^4.1.5",
 		"prettier": "^3.4.2",
+		"rimraf": "^6.0.1",
 		"ts-jest": "^29.2.5",
 		"typescript": "^5.4.5"
 	},

package.json

@@ -0,0 +1,6 @@
+module.exports = async function getFolderSize() {
+	return {
+		size: 1000,
+		errors: [],
+	}
+}

src/__mocks__/get-folder-size.js

@@ -6,13 +6,14 @@ import delay from "delay"
 import fs from "fs/promises"
 import os from "os"
 import pWaitFor from "p-wait-for"
+import getFolderSize from "get-folder-size"
 import * as path from "path"
 import { serializeError } from "serialize-error"
 import * as vscode from "vscode"
 import { ApiHandler, SingleCompletionHandler, buildApiHandler } from "../api"
 import { ApiStream } from "../api/transform/stream"
 import { DIFF_VIEW_URI_SCHEME, DiffViewProvider } from "../integrations/editor/DiffViewProvider"
-import { CheckpointService } from "../services/checkpoints/CheckpointService"
+import { CheckpointService, CheckpointServiceFactory } from "../services/checkpoints"
 import { findToolName, formatContentBlockToMarkdown } from "../integrations/misc/export-markdown"
 import {
 	extractTextFromFile,
@@ -239,7 +240,8 @@ export class Cline {
 
 	private async saveClineMessages() {
 		try {
-			const filePath = path.join(await this.ensureTaskDirectoryExists(), GlobalFileNames.uiMessages)
+			const taskDir = await this.ensureTaskDirectoryExists()
+			const filePath = path.join(taskDir, GlobalFileNames.uiMessages)
 			await fs.writeFile(filePath, JSON.stringify(this.clineMessages))
 			// combined as they are in ChatView
 			const apiMetrics = getApiMetrics(combineApiRequests(combineCommandSequences(this.clineMessages.slice(1))))
@@ -251,6 +253,17 @@ export class Cline {
 						(m) => !(m.ask === "resume_task" || m.ask === "resume_completed_task"),
 					)
 				]
+
+			let taskDirSize = 0
+
+			try {
+				taskDirSize = await getFolderSize.loose(taskDir)
+			} catch (err) {
+				console.error(
+					`[saveClineMessages] failed to get task directory size (${taskDir}): ${err instanceof Error ? err.message : String(err)}`,
+				)
+			}
+
 			await this.providerRef.deref()?.updateTaskHistory({
 				id: this.taskId,
 				ts: lastRelevantMessage.ts,
@@ -260,6 +273,7 @@ export class Cline {
 				cacheWrites: apiMetrics.totalCacheWrites,
 				cacheReads: apiMetrics.totalCacheReads,
 				totalCost: apiMetrics.totalCost,
+				size: taskDirSize,
 			})
 		} catch (error) {
 			console.error("Failed to save cline messages:", error)
@@ -2692,7 +2706,7 @@ export class Cline {
 		}
 
 		if (isCheckpointPossible) {
-			await this.checkpointSave()
+			await this.checkpointSave({ isFirst: false })
 		}
 
 		/*
@@ -2762,7 +2776,7 @@ export class Cline {
 		const isFirstRequest = this.clineMessages.filter((m) => m.say === "api_req_started").length === 0
 
 		if (isFirstRequest) {
-			await this.checkpointSave()
+			await this.checkpointSave({ isFirst: true })
 		}
 
 		// 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
@@ -3255,11 +3269,32 @@ export class Cline {
 	// Checkpoints
 
 	private async getCheckpointService() {
+		if (!this.checkpointsEnabled) {
+			throw new Error("Checkpoints are disabled")
+		}
+
 		if (!this.checkpointService) {
-			this.checkpointService = await CheckpointService.create({
-				taskId: this.taskId,
-				baseDir: vscode.workspace.workspaceFolders?.map((folder) => folder.uri.fsPath).at(0) ?? "",
-				log: (message) => this.providerRef.deref()?.log(message),
+			const workspaceDir = vscode.workspace.workspaceFolders?.map((folder) => folder.uri.fsPath).at(0)
+			const shadowDir = this.providerRef.deref()?.context.globalStorageUri.fsPath
+
+			if (!workspaceDir) {
+				this.providerRef.deref()?.log("[getCheckpointService] workspace folder not found")
+				throw new Error("Workspace directory not found")
+			}
+
+			if (!shadowDir) {
+				this.providerRef.deref()?.log("[getCheckpointService] shadowDir not found")
+				throw new Error("Global storage directory not found")
+			}
+
+			this.checkpointService = await CheckpointServiceFactory.create({
+				strategy: "shadow",
+				options: {
+					taskId: this.taskId,
+					workspaceDir,
+					shadowDir,
+					log: (message) => this.providerRef.deref()?.log(message),
+				},
 			})
 		}
 
@@ -3318,29 +3353,25 @@ export class Cline {
 		}
 	}
 
-	public async checkpointSave() {
+	public async checkpointSave({ isFirst }: { isFirst: boolean }) {
 		if (!this.checkpointsEnabled) {
 			return
 		}
 
 		try {
-			const isFirst = !this.checkpointService
 			const service = await this.getCheckpointService()
-			const commit = await service.saveCheckpoint(`Task: ${this.taskId}, Time: ${Date.now()}`)
+			const strategy = service.strategy
+			const version = service.version
 
-			if (commit?.commit) {
-				await this.providerRef
-					.deref()
-					?.postMessageToWebview({ type: "currentCheckpointUpdated", text: service.currentCheckpoint })
+			const commit = await service.saveCheckpoint(`Task: ${this.taskId}, Time: ${Date.now()}`)
+			const fromHash = service.baseHash
+			const toHash = isFirst ? commit?.commit || fromHash : commit?.commit
 
-				// Checkpoint metadata required by the UI.
-				const checkpoint = {
-					isFirst,
-					from: service.baseCommitHash,
-					to: commit.commit,
-				}
+			if (toHash) {
+				await this.providerRef.deref()?.postMessageToWebview({ type: "currentCheckpointUpdated", text: toHash })
 
-				await this.say("checkpoint_saved", commit.commit, undefined, undefined, checkpoint)
+				const checkpoint = { isFirst, from: fromHash, to: toHash, strategy, version }
+				await this.say("checkpoint_saved", toHash, undefined, undefined, checkpoint)
 			}
 		} catch (err) {
 			this.providerRef.deref()?.log("[checkpointSave] disabling checkpoints for this task")
@@ -3371,9 +3402,7 @@ export class Cline {
 			const service = await this.getCheckpointService()
 			await service.restoreCheckpoint(commitHash)
 
-			await this.providerRef
-				.deref()
-				?.postMessageToWebview({ type: "currentCheckpointUpdated", text: service.currentCheckpoint })
+			await this.providerRef.deref()?.postMessageToWebview({ type: "currentCheckpointUpdated", text: commitHash })
 
 			if (mode === "restore") {
 				await this.overwriteApiConversationHistory(

src/core/Cline.ts

@@ -2277,35 +2277,55 @@ export class ClineProvider implements vscode.WebviewViewProvider {
 
 		await this.deleteTaskFromState(id)
 
-		// Delete the task files
+		// Delete the task files.
 		const apiConversationHistoryFileExists = await fileExistsAtPath(apiConversationHistoryFilePath)
+
 		if (apiConversationHistoryFileExists) {
 			await fs.unlink(apiConversationHistoryFilePath)
 		}
+
 		const uiMessagesFileExists = await fileExistsAtPath(uiMessagesFilePath)
+
 		if (uiMessagesFileExists) {
 			await fs.unlink(uiMessagesFilePath)
 		}
+
 		const legacyMessagesFilePath = path.join(taskDirPath, "claude_messages.json")
+
 		if (await fileExistsAtPath(legacyMessagesFilePath)) {
 			await fs.unlink(legacyMessagesFilePath)
 		}
-		await fs.rmdir(taskDirPath) // succeeds if the dir is empty
 
 		const { checkpointsEnabled } = await this.getState()
 		const baseDir = vscode.workspace.workspaceFolders?.map((folder) => folder.uri.fsPath).at(0)
-		const branch = `roo-code-checkpoints-${id}`
 
+		// Delete checkpoints branch.
 		if (checkpointsEnabled && baseDir) {
+			const branchSummary = await simpleGit(baseDir)
+				.branch(["-D", `roo-code-checkpoints-${id}`])
+				.catch(() => undefined)
+
+			if (branchSummary) {
+				console.log(`[deleteTaskWithId${id}] deleted checkpoints branch`)
+			}
+		}
+
+		// Delete checkpoints directory
+		const checkpointsDir = path.join(taskDirPath, "checkpoints")
+
+		if (await fileExistsAtPath(checkpointsDir)) {
 			try {
-				await simpleGit(baseDir).branch(["-D", branch])
-				console.log(`[deleteTaskWithId] Deleted branch ${branch}`)
-			} catch (err) {
+				await fs.rm(checkpointsDir, { recursive: true, force: true })
+				console.log(`[deleteTaskWithId${id}] removed checkpoints repo`)
+			} catch (error) {
 				console.error(
-					`[deleteTaskWithId] Error deleting branch ${branch}: ${err instanceof Error ? err.message : String(err)}`,
+					`[deleteTaskWithId${id}] failed to remove checkpoints repo: ${error instanceof Error ? error.message : String(error)}`,
 				)
 			}
 		}
+
+		// Succeeds if the dir is empty.
+		await fs.rmdir(taskDirPath)
 	}
 
 	async deleteTaskFromState(id: string) {
@@ -2373,6 +2393,9 @@ export class ClineProvider implements vscode.WebviewViewProvider {
 			alwaysAllowMcp: alwaysAllowMcp ?? false,
 			alwaysAllowModeSwitch: alwaysAllowModeSwitch ?? false,
 			uriScheme: vscode.env.uriScheme,
+			currentTaskItem: this.cline?.taskId
+				? (taskHistory || []).find((item) => item.id === this.cline?.taskId)
+				: undefined,
 			clineMessages: this.cline?.clineMessages || [],
 			taskHistory: (taskHistory || [])
 				.filter((item: HistoryItem) => item.ts && item.task)

src/core/webview/ClineProvider.ts

@@ -0,0 +1,29 @@
+import { LocalCheckpointService, LocalCheckpointServiceOptions } from "./LocalCheckpointService"
+import { ShadowCheckpointService, ShadowCheckpointServiceOptions } from "./ShadowCheckpointService"
+
+export type CreateCheckpointServiceFactoryOptions =
+	| {
+			strategy: "local"
+			options: LocalCheckpointServiceOptions
+	  }
+	| {
+			strategy: "shadow"
+			options: ShadowCheckpointServiceOptions
+	  }
+
+type CheckpointServiceType<T extends CreateCheckpointServiceFactoryOptions> = T extends { strategy: "local" }
+	? LocalCheckpointService
+	: T extends { strategy: "shadow" }
+		? ShadowCheckpointService
+		: never
+
+export class CheckpointServiceFactory {
+	public static create<T extends CreateCheckpointServiceFactoryOptions>(options: T): CheckpointServiceType<T> {
+		switch (options.strategy) {
+			case "local":
+				return LocalCheckpointService.create(options.options) as any
+			case "shadow":
+				return ShadowCheckpointService.create(options.options) as any
+		}
+	}
+}