Await checkpoint saves (except the initial) (RooCodeInc#2665)

Author: cte

src/core/Cline.ts

@@ -16,11 +16,7 @@ import { TokenUsage } from "../schemas"
 import { ApiHandler, buildApiHandler } from "../api"
 import { ApiStream } from "../api/transform/stream"
 import { DIFF_VIEW_URI_SCHEME, DiffViewProvider } from "../integrations/editor/DiffViewProvider"
-import {
-	CheckpointServiceOptions,
-	RepoPerTaskCheckpointService,
-	RepoPerWorkspaceCheckpointService,
-} from "../services/checkpoints"
+import { CheckpointServiceOptions, RepoPerTaskCheckpointService } from "../services/checkpoints"
 import { findToolName, formatContentBlockToMarkdown } from "../integrations/misc/export-markdown"
 import { fetchInstructionsTool } from "./tools/fetchInstructionsTool"
 import { listFilesTool } from "./tools/listFilesTool"
@@ -30,7 +26,6 @@ import { Terminal } from "../integrations/terminal/Terminal"
 import { TerminalRegistry } from "../integrations/terminal/TerminalRegistry"
 import { UrlContentFetcher } from "../services/browser/UrlContentFetcher"
 import { listFiles } from "../services/glob/list-files"
-import { CheckpointStorage } from "../shared/checkpoints"
 import { ApiConfiguration } from "../shared/api"
 import { findLastIndex } from "../shared/array"
 import { combineApiRequests } from "../shared/combineApiRequests"
@@ -104,7 +99,6 @@ export type ClineOptions = {
 	customInstructions?: string
 	enableDiff?: boolean
 	enableCheckpoints?: boolean
-	checkpointStorage?: CheckpointStorage
 	fuzzyMatchThreshold?: number
 	consecutiveMistakeLimit?: number
 	task?: string
@@ -162,8 +156,8 @@ export class Cline extends EventEmitter<ClineEvents> {
 
 	// checkpoints
 	private enableCheckpoints: boolean
-	private checkpointStorage: CheckpointStorage
-	private checkpointService?: RepoPerTaskCheckpointService | RepoPerWorkspaceCheckpointService
+	private checkpointService?: RepoPerTaskCheckpointService
+	private checkpointServiceInitializing = false
 
 	// streaming
 	isWaitingForFirstChunk = false
@@ -184,7 +178,6 @@ export class Cline extends EventEmitter<ClineEvents> {
 		customInstructions,
 		enableDiff = false,
 		enableCheckpoints = true,
-		checkpointStorage = "task",
 		fuzzyMatchThreshold = 1.0,
 		consecutiveMistakeLimit = 3,
 		task,
@@ -223,7 +216,6 @@ export class Cline extends EventEmitter<ClineEvents> {
 		this.providerRef = new WeakRef(provider)
 		this.diffViewProvider = new DiffViewProvider(this.cwd)
 		this.enableCheckpoints = enableCheckpoints
-		this.checkpointStorage = checkpointStorage
 
 		this.rootTask = rootTask
 		this.parentTask = parentTask
@@ -1680,9 +1672,11 @@ export class Cline extends EventEmitter<ClineEvents> {
 		}
 
 		const recentlyModifiedFiles = this.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
-			this.checkpointSave()
+			// TODO: We can track what file changes were made and only
+			// checkpoint those files, this will be save storage.
+			await this.checkpointSave()
 		}
 
 		/*
@@ -2397,6 +2391,11 @@ export class Cline extends EventEmitter<ClineEvents> {
 			return this.checkpointService
 		}
 
+		if (this.checkpointServiceInitializing) {
+			console.log("[Cline#getCheckpointService] checkpoint service is still initializing")
+			return undefined
+		}
+
 		const log = (message: string) => {
 			console.log(message)
 
@@ -2407,19 +2406,21 @@ export class Cline extends EventEmitter<ClineEvents> {
 			}
 		}
 
+		console.log("[Cline#getCheckpointService] initializing checkpoints service")
+
 		try {
 			const workspaceDir = getWorkspacePath()
 
 			if (!workspaceDir) {
-				log("[Cline#initializeCheckpoints] workspace folder not found, disabling checkpoints")
+				log("[Cline#getCheckpointService] workspace folder not found, disabling checkpoints")
 				this.enableCheckpoints = false
 				return undefined
 			}
 
 			const globalStorageDir = this.providerRef.deref()?.context.globalStorageUri.fsPath
 
 			if (!globalStorageDir) {
-				log("[Cline#initializeCheckpoints] globalStorageDir not found, disabling checkpoints")
+				log("[Cline#getCheckpointService] globalStorageDir not found, disabling checkpoints")
 				this.enableCheckpoints = false
 				return undefined
 			}
@@ -2431,28 +2432,26 @@ export class Cline extends EventEmitter<ClineEvents> {
 				log,
 			}
 
-			// Only `task` is supported at the moment until we figure out how
-			// to fully isolate the `workspace` variant.
-			// const service =
-			// 	this.checkpointStorage === "task"
-			// 		? RepoPerTaskCheckpointService.create(options)
-			// 		: RepoPerWorkspaceCheckpointService.create(options)
-
 			const service = RepoPerTaskCheckpointService.create(options)
 
+			this.checkpointServiceInitializing = true
+
 			service.on("initialize", () => {
+				log("[Cline#getCheckpointService] service initialized")
+
 				try {
 					const isCheckpointNeeded =
 						typeof this.clineMessages.find(({ say }) => say === "checkpoint_saved") === "undefined"
 
 					this.checkpointService = service
+					this.checkpointServiceInitializing = false
 
 					if (isCheckpointNeeded) {
-						log("[Cline#initializeCheckpoints] no checkpoints found, saving initial checkpoint")
+						log("[Cline#getCheckpointService] no checkpoints found, saving initial checkpoint")
 						this.checkpointSave()
 					}
 				} catch (err) {
-					log("[Cline#initializeCheckpoints] caught error in on('initialize'), disabling checkpoints")
+					log("[Cline#getCheckpointService] caught error in on('initialize'), disabling checkpoints")
 					this.enableCheckpoints = false
 				}
 			})
@@ -2462,29 +2461,31 @@ export class Cline extends EventEmitter<ClineEvents> {
 					this.providerRef.deref()?.postMessageToWebview({ type: "currentCheckpointUpdated", text: to })
 
 					this.say("checkpoint_saved", to, undefined, undefined, { isFirst, from, to }).catch((err) => {
-						log("[Cline#initializeCheckpoints] caught unexpected error in say('checkpoint_saved')")
+						log("[Cline#getCheckpointService] caught unexpected error in say('checkpoint_saved')")
 						console.error(err)
 					})
 				} catch (err) {
 					log(
-						"[Cline#initializeCheckpoints] caught unexpected error in on('checkpoint'), disabling checkpoints",
+						"[Cline#getCheckpointService] caught unexpected error in on('checkpoint'), disabling checkpoints",
 					)
 					console.error(err)
 					this.enableCheckpoints = false
 				}
 			})
 
+			log("[Cline#getCheckpointService] initializing shadow git")
+
 			service.initShadowGit().catch((err) => {
 				log(
-					`[Cline#initializeCheckpoints] caught unexpected error in initShadowGit, disabling checkpoints (${err.message})`,
+					`[Cline#getCheckpointService] caught unexpected error in initShadowGit, disabling checkpoints (${err.message})`,
 				)
 				console.error(err)
 				this.enableCheckpoints = false
 			})
 
 			return service
 		} catch (err) {
-			log("[Cline#initializeCheckpoints] caught unexpected error, disabling checkpoints")
+			log("[Cline#getCheckpointService] caught unexpected error, disabling checkpoints")
 			this.enableCheckpoints = false
 			return undefined
 		}
@@ -2508,6 +2509,7 @@ export class Cline extends EventEmitter<ClineEvents> {
 				},
 				{ interval, timeout },
 			)
+
 			return service
 		} catch (err) {
 			return undefined
@@ -2569,7 +2571,7 @@ export class Cline extends EventEmitter<ClineEvents> {
 		}
 	}
 
-	public checkpointSave() {
+	public async checkpointSave() {
 		const service = this.getCheckpointService()
 
 		if (!service) {
@@ -2580,14 +2582,15 @@ export class Cline extends EventEmitter<ClineEvents> {
 			this.providerRef
 				.deref()
 				?.log("[checkpointSave] checkpoints didn't initialize in time, disabling checkpoints for this task")
+
 			this.enableCheckpoints = false
 			return
 		}
 
 		telemetryService.captureCheckpointCreated(this.taskId)
 
 		// Start the checkpoint process in the background.
-		service.saveCheckpoint(`Task: ${this.taskId}, Time: ${Date.now()}`).catch((err) => {
+		return service.saveCheckpoint(`Task: ${this.taskId}, Time: ${Date.now()}`).catch((err) => {
 			console.error("[Cline#checkpointSave] caught unexpected error, disabling checkpoints", err)
 			this.enableCheckpoints = false
 		})

src/core/webview/ClineProvider.ts

@@ -483,7 +483,6 @@ export class ClineProvider extends EventEmitter<ClineProviderEvents> implements
 				| "customInstructions"
 				| "enableDiff"
 				| "enableCheckpoints"
-				| "checkpointStorage"
 				| "fuzzyMatchThreshold"
 				| "consecutiveMistakeLimit"
 				| "experiments"
@@ -495,7 +494,6 @@ export class ClineProvider extends EventEmitter<ClineProviderEvents> implements
 			customModePrompts,
 			diffEnabled: enableDiff,
 			enableCheckpoints,
-			checkpointStorage,
 			fuzzyMatchThreshold,
 			mode,
 			customInstructions: globalInstructions,
@@ -511,7 +509,6 @@ export class ClineProvider extends EventEmitter<ClineProviderEvents> implements
 			customInstructions: effectiveInstructions,
 			enableDiff,
 			enableCheckpoints,
-			checkpointStorage,
 			fuzzyMatchThreshold,
 			task,
 			images,
@@ -540,7 +537,6 @@ export class ClineProvider extends EventEmitter<ClineProviderEvents> implements
 			customModePrompts,
 			diffEnabled: enableDiff,
 			enableCheckpoints,
-			checkpointStorage,
 			fuzzyMatchThreshold,
 			mode,
 			customInstructions: globalInstructions,
@@ -550,38 +546,12 @@ export class ClineProvider extends EventEmitter<ClineProviderEvents> implements
 		const modePrompt = customModePrompts?.[mode] as PromptComponent
 		const effectiveInstructions = [globalInstructions, modePrompt?.customInstructions].filter(Boolean).join("\n\n")
 
-		const taskId = historyItem.id
-		const globalStorageDir = this.contextProxy.globalStorageUri.fsPath
-		const workspaceDir = this.cwd
-
-		const checkpoints: Pick<ClineOptions, "enableCheckpoints" | "checkpointStorage"> = {
-			enableCheckpoints,
-			checkpointStorage,
-		}
-
-		if (enableCheckpoints) {
-			try {
-				checkpoints.checkpointStorage = await ShadowCheckpointService.getTaskStorage({
-					taskId,
-					globalStorageDir,
-					workspaceDir,
-				})
-
-				this.log(
-					`[ClineProvider#initClineWithHistoryItem] Using ${checkpoints.checkpointStorage} storage for ${taskId}`,
-				)
-			} catch (error) {
-				checkpoints.enableCheckpoints = false
-				this.log(`[ClineProvider#initClineWithHistoryItem] Error getting task storage: ${error.message}`)
-			}
-		}
-
 		const cline = new Cline({
 			provider: this,
 			apiConfiguration,
 			customInstructions: effectiveInstructions,
 			enableDiff,
-			...checkpoints,
+			enableCheckpoints,
 			fuzzyMatchThreshold,
 			historyItem,
 			experiments,
@@ -1210,7 +1180,6 @@ export class ClineProvider extends EventEmitter<ClineProviderEvents> implements
 			ttsSpeed,
 			diffEnabled,
 			enableCheckpoints,
-			checkpointStorage,
 			taskHistory,
 			soundVolume,
 			browserViewportSize,
@@ -1282,7 +1251,6 @@ export class ClineProvider extends EventEmitter<ClineProviderEvents> implements
 			ttsSpeed: ttsSpeed ?? 1.0,
 			diffEnabled: diffEnabled ?? true,
 			enableCheckpoints: enableCheckpoints ?? true,
-			checkpointStorage: checkpointStorage ?? "task",
 			shouldShowAnnouncement:
 				telemetrySetting !== "unset" && lastShownAnnouncementId !== this.latestAnnouncementId,
 			allowedCommands,
@@ -1377,7 +1345,6 @@ export class ClineProvider extends EventEmitter<ClineProviderEvents> implements
 			ttsSpeed: stateValues.ttsSpeed ?? 1.0,
 			diffEnabled: stateValues.diffEnabled ?? true,
 			enableCheckpoints: stateValues.enableCheckpoints ?? true,
-			checkpointStorage: stateValues.checkpointStorage ?? "task",
 			soundVolume: stateValues.soundVolume,
 			browserViewportSize: stateValues.browserViewportSize ?? "900x600",
 			screenshotQuality: stateValues.screenshotQuality ?? 75,

src/core/webview/__tests__/ClineProvider.test.ts

@@ -407,7 +407,6 @@ describe("ClineProvider", () => {
 			ttsEnabled: false,
 			diffEnabled: false,
 			enableCheckpoints: false,
-			checkpointStorage: "task",
 			writeDelayMs: 1000,
 			browserViewportSize: "900x600",
 			fuzzyMatchThreshold: 1.0,
@@ -829,7 +828,6 @@ describe("ClineProvider", () => {
 			mode: "code",
 			diffEnabled: true,
 			enableCheckpoints: false,
-			checkpointStorage: "task",
 			fuzzyMatchThreshold: 1.0,
 			experiments: experimentDefault,
 		} as any)
@@ -848,7 +846,6 @@ describe("ClineProvider", () => {
 			customInstructions: modeCustomInstructions,
 			enableDiff: true,
 			enableCheckpoints: false,
-			checkpointStorage: "task",
 			fuzzyMatchThreshold: 1.0,
 			task: "Test task",
 			experiments: experimentDefault,

src/core/webview/webviewMessageHandler.ts

@@ -4,7 +4,7 @@ import pWaitFor from "p-wait-for"
 import * as vscode from "vscode"
 
 import { ClineProvider } from "./ClineProvider"
-import { CheckpointStorage, Language, ApiConfigMeta } from "../../schemas"
+import { Language, ApiConfigMeta } from "../../schemas"
 import { changeLanguage, t } from "../../i18n"
 import { ApiConfiguration } from "../../shared/api"
 import { supportPrompt } from "../../shared/support-prompt"
@@ -655,12 +655,6 @@ export const webviewMessageHandler = async (provider: ClineProvider, message: We
 			await updateGlobalState("enableCheckpoints", enableCheckpoints)
 			await provider.postStateToWebview()
 			break
-		case "checkpointStorage":
-			console.log(`[ClineProvider] checkpointStorage: ${message.text}`)
-			const checkpointStorage = message.text ?? "task"
-			await updateGlobalState("checkpointStorage", checkpointStorage as CheckpointStorage)
-			await provider.postStateToWebview()
-			break
 		case "browserViewportSize":
 			const browserViewportSize = message.text ?? "900x600"
 			await updateGlobalState("browserViewportSize", browserViewportSize)

src/exports/roo-code.d.ts

@@ -259,7 +259,6 @@ type GlobalSettings = {
 	remoteBrowserHost?: string | undefined
 	cachedChromeHostUrl?: string | undefined
 	enableCheckpoints?: boolean | undefined
-	checkpointStorage?: ("task" | "workspace") | undefined
 	showGreeting?: boolean | undefined
 	ttsEnabled?: boolean | undefined
 	ttsSpeed?: number | undefined

src/exports/types.ts

@@ -262,7 +262,6 @@ type GlobalSettings = {
 	remoteBrowserHost?: string | undefined
 	cachedChromeHostUrl?: string | undefined
 	enableCheckpoints?: boolean | undefined
-	checkpointStorage?: ("task" | "workspace") | undefined
 	showGreeting?: boolean | undefined
 	ttsEnabled?: boolean | undefined
 	ttsSpeed?: number | undefined