From 0bd56ad165a61bb90b6546ad3e0e0b2c906c84c3 Mon Sep 17 00:00:00 2001 From: cte Date: Tue, 4 Mar 2025 01:37:05 -0800 Subject: [PATCH 1/2] Checkpoints: one-branch-per-task / one-repo-per-workspace --- src/core/Cline.ts | 162 +++++++++----- src/core/webview/ClineProvider.ts | 4 + .../RepoPerTaskCheckpointService.ts | 15 ++ .../RepoPerWorkspaceCheckpointService.ts | 76 +++++++ .../checkpoints/ShadowCheckpointService.ts | 211 ++++++++++-------- .../__tests__/ShadowCheckpointService.test.ts | 117 +++++----- src/services/checkpoints/index.ts | 4 + src/services/checkpoints/types.ts | 36 +-- src/shared/WebviewMessage.ts | 1 + .../chat/checkpoints/CheckpointMenu.tsx | 10 +- .../src/components/history/HistoryView.tsx | 7 +- 11 files changed, 407 insertions(+), 236 deletions(-) create mode 100644 src/services/checkpoints/RepoPerTaskCheckpointService.ts create mode 100644 src/services/checkpoints/RepoPerWorkspaceCheckpointService.ts create mode 100644 src/services/checkpoints/index.ts diff --git a/src/core/Cline.ts b/src/core/Cline.ts index b9b243dbc1..bccac0575f 100644 --- a/src/core/Cline.ts +++ b/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 = "task", 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 { - 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,47 +3381,45 @@ 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 }) => { + const service = + this.checkpointStorage === "task" + ? RepoPerTaskCheckpointService.create(options) + : RepoPerWorkspaceCheckpointService.create(options) + + service.on("initialize", () => { 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 isCheckpointNeeded = + typeof this.clineMessages.find(({ say }) => say === "checkpoint_saved") === "undefined" this.checkpointService = service - this.checkpointSave() + + if (isCheckpointNeeded) { + log("[Cline#initializeCheckpoints] no checkpoints found, saving initial checkpoint") + this.checkpointSave() + } } catch (err) { log("[Cline#initializeCheckpoints] caught error in on('initialize'), disabling checkpoints") this.enableCheckpoints = false @@ -3417,41 +3428,77 @@ 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) => - console.error("Error saving checkpoint message:", e), - ) + this.say("checkpoint_saved", to, undefined, undefined, { isFirst, from, to }).catch((err) => { + log("[Cline#initializeCheckpoints] caught unexpected error in say('checkpoint_saved')") + console.error(err) + }) } catch (err) { - log("[Cline#initializeCheckpoints] caught error in on('checkpoint'), disabling checkpoints") + log( + "[Cline#initializeCheckpoints] caught unexpected error in on('checkpoint'), disabling checkpoints", + ) + console.error(err) this.enableCheckpoints = false } }) - service.initShadowGit() + service.initShadowGit().catch((err) => { + log("[Cline#initializeCheckpoints] caught unexpected error in initShadowGit, disabling checkpoints") + console.error(err) + this.enableCheckpoints = false + }) + + 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 } } public async checkpointDiff({ ts, + previousCommitHash, commitHash, mode, }: { ts: number + previousCommitHash?: string commitHash: string mode: "full" | "checkpoint" }) { - if (!this.checkpointService || !this.enableCheckpoints) { + const service = await this.getInitializedCheckpointService() + + if (!service) { return } - let previousCommitHash = undefined - - if (mode === "checkpoint") { + if (!previousCommitHash && mode === "checkpoint") { const previousCheckpoint = this.clineMessages .filter(({ say }) => say === "checkpoint_saved") .sort((a, b) => b.ts - a.ts) @@ -3461,7 +3508,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 +3535,25 @@ 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()}`).catch((err) => { + console.error("[Cline#checkpointSave] caught unexpected error, disabling checkpoints", err) + this.enableCheckpoints = false + }) } public async checkpointRestore({ @@ -3505,7 +3565,9 @@ export class Cline { commitHash: string mode: "preview" | "restore" }) { - if (!this.checkpointService || !this.enableCheckpoints) { + const service = await this.getInitializedCheckpointService() + + if (!service) { return } @@ -3516,7 +3578,7 @@ export class Cline { } try { - await this.checkpointService.restoreCheckpoint(commitHash) + await service.restoreCheckpoint(commitHash) await this.providerRef.deref()?.postMessageToWebview({ type: "currentCheckpointUpdated", text: commitHash }) diff --git a/src/core/webview/ClineProvider.ts b/src/core/webview/ClineProvider.ts index 0a2e385b6a..d318989dfa 100644 --- a/src/core/webview/ClineProvider.ts +++ b/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 diff --git a/src/services/checkpoints/RepoPerTaskCheckpointService.ts b/src/services/checkpoints/RepoPerTaskCheckpointService.ts new file mode 100644 index 0000000000..2190ed302d --- /dev/null +++ b/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, + ) + } +} diff --git a/src/services/checkpoints/RepoPerWorkspaceCheckpointService.ts b/src/services/checkpoints/RepoPerWorkspaceCheckpointService.ts new file mode 100644 index 0000000000..654e90f1c3 --- /dev/null +++ b/src/services/checkpoints/RepoPerWorkspaceCheckpointService.ts @@ -0,0 +1,76 @@ +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(source: string) { + if (!this.git) { + throw new Error("Shadow git repo not initialized") + } + + const startTime = Date.now() + const branch = `roo-${this.taskId}` + const currentBranch = await this.git.revparse(["--abbrev-ref", "HEAD"]) + + if (currentBranch === branch) { + return + } + + this.log(`[${this.constructor.name}#checkoutTaskBranch{${source}}] checking out ${branch}`) + const branches = await this.git.branchLocal() + let exists = branches.all.includes(branch) + + if (!exists) { + await this.git.checkoutLocalBranch(branch) + } else { + await this.git.checkout(branch) + } + + const duration = Date.now() - startTime + + this.log( + `[${this.constructor.name}#checkoutTaskBranch{${source}}] ${exists ? "checked out" : "created"} branch "${branch}" in ${duration}ms`, + ) + } + + override async initShadowGit() { + return await super.initShadowGit(() => this.checkoutTaskBranch("initShadowGit")) + } + + override async saveCheckpoint(message: string) { + await this.checkoutTaskBranch("saveCheckpoint") + return super.saveCheckpoint(message) + } + + override async restoreCheckpoint(commitHash: string) { + await this.checkoutTaskBranch("restoreCheckpoint") + await super.restoreCheckpoint(commitHash) + } + + override async getDiff({ from, to }: { from?: string; to?: string }) { + if (!this.git) { + throw new Error("Shadow git repo not initialized") + } + + await this.checkoutTaskBranch("getDiff") + + if (!from && to) { + from = `${to}~` + } + + return super.getDiff({ from, to }) + } + + public static create({ taskId, workspaceDir, shadowDir, log = console.log }: CheckpointServiceOptions) { + const workspaceHash = crypto.createHash("sha256").update(workspaceDir).digest("hex").toString().slice(0, 8) + + return new RepoPerWorkspaceCheckpointService( + taskId, + path.join(shadowDir, "checkpoints", workspaceHash), + workspaceDir, + log, + ) + } +} diff --git a/src/services/checkpoints/ShadowCheckpointService.ts b/src/services/checkpoints/ShadowCheckpointService.ts index bdad6d686a..54e12ce753 100644 --- a/src/services/checkpoints/ShadowCheckpointService.ts +++ b/src/services/checkpoints/ShadowCheckpointService.ts @@ -1,58 +1,70 @@ import fs from "fs/promises" import os from "os" import * as path from "path" -import { globby } from "globby" +import EventEmitter from "events" + import simpleGit, { SimpleGit } from "simple-git" +import { globby } from "globby" import { GIT_DISABLED_SUFFIX, GIT_EXCLUDES } from "./constants" -import { CheckpointService, CheckpointServiceOptions, CheckpointEventEmitter } from "./types" +import { CheckpointDiff, CheckpointResult, CheckpointEventMap } from "./types" -export interface ShadowCheckpointServiceOptions extends CheckpointServiceOptions { - shadowDir: string -} +export abstract class ShadowCheckpointService extends EventEmitter { + public readonly taskId: string + public readonly checkpointsDir: string + public readonly workspaceDir: string -export class ShadowCheckpointService extends CheckpointEventEmitter implements CheckpointService { - public readonly version = 1 + protected _checkpoints: string[] = [] + protected _baseHash?: string - private _checkpoints: string[] = [] - private _baseHash?: string - private _isInitialized = false + protected readonly dotGitDir: string + protected git?: SimpleGit + protected readonly log: (message: string) => void + protected shadowGitConfigWorktree?: string public get baseHash() { return this._baseHash } - private set baseHash(value: string | undefined) { + protected set baseHash(value: string | undefined) { this._baseHash = value } public get isInitialized() { - return this._isInitialized + return !!this.git } - private set isInitialized(value: boolean) { - this._isInitialized = value - } - - private readonly shadowGitDir: string - private shadowGitConfigWorktree?: string - - private constructor( - public readonly taskId: string, - public readonly git: SimpleGit, - public readonly shadowDir: string, - public readonly workspaceDir: string, - private readonly log: (message: string) => void, - ) { + constructor(taskId: string, checkpointsDir: string, workspaceDir: string, log: (message: string) => void) { super() - this.shadowGitDir = path.join(this.shadowDir, "tasks", this.taskId, "checkpoints", ".git") + + const homedir = os.homedir() + const desktopPath = path.join(homedir, "Desktop") + const documentsPath = path.join(homedir, "Documents") + const downloadsPath = path.join(homedir, "Downloads") + const protectedPaths = [homedir, desktopPath, documentsPath, downloadsPath] + + if (protectedPaths.includes(workspaceDir)) { + throw new Error(`Cannot use checkpoints in ${workspaceDir}`) + } + + this.taskId = taskId + this.checkpointsDir = checkpointsDir + this.workspaceDir = workspaceDir + + this.dotGitDir = path.join(this.checkpointsDir, ".git") + this.log = log } - public async initShadowGit() { - if (this.isInitialized) { - return + public async initShadowGit(onInit?: () => Promise) { + if (this.git) { + throw new Error("Shadow git repo already initialized") } + await fs.mkdir(this.checkpointsDir, { recursive: true }) + const git = simpleGit(this.checkpointsDir) + const gitVersion = await git.version() + this.log(`[${this.constructor.name}#create] git = ${gitVersion}`) + const fileExistsAtPath = (path: string) => fs .access(path) @@ -62,9 +74,9 @@ export class ShadowCheckpointService extends CheckpointEventEmitter implements C let created = false const startTime = Date.now() - if (await fileExistsAtPath(this.shadowGitDir)) { - this.log(`[CheckpointService#initShadowGit] shadow git repo already exists at ${this.shadowGitDir}`) - const worktree = await this.getShadowGitConfigWorktree() + if (await fileExistsAtPath(this.dotGitDir)) { + this.log(`[${this.constructor.name}#initShadowGit] shadow git repo already exists at ${this.dotGitDir}`) + const worktree = await this.getShadowGitConfigWorktree(git) if (worktree !== this.workspaceDir) { throw new Error( @@ -72,15 +84,15 @@ export class ShadowCheckpointService extends CheckpointEventEmitter implements C ) } - this.baseHash = await this.git.revparse(["--abbrev-ref", "HEAD"]) + this.baseHash = await git.revparse(["HEAD"]) } else { - this.log(`[CheckpointService#initShadowGit] creating shadow git repo at ${this.workspaceDir}`) + this.log(`[${this.constructor.name}#initShadowGit] creating shadow git repo at ${this.checkpointsDir}`) - await this.git.init() - await this.git.addConfig("core.worktree", this.workspaceDir) // Sets the working tree to the current workspace. - await this.git.addConfig("commit.gpgSign", "false") // Disable commit signing for shadow repo. - await this.git.addConfig("user.name", "Roo Code") - await this.git.addConfig("user.email", "noreply@example.com") + await git.init() + await git.addConfig("core.worktree", this.workspaceDir) // Sets the working tree to the current workspace. + await git.addConfig("commit.gpgSign", "false") // Disable commit signing for shadow repo. + await git.addConfig("user.name", "Roo Code") + await git.addConfig("user.email", "noreply@example.com") let lfsPatterns: string[] = [] // Get LFS patterns from workspace if they exist. @@ -95,7 +107,7 @@ export class ShadowCheckpointService extends CheckpointEventEmitter implements C } } catch (error) { this.log( - `[CheckpointService#initShadowGit] failed to read .gitattributes: ${error instanceof Error ? error.message : String(error)}`, + `[${this.constructor.name}#initShadowGit] failed to read .gitattributes: ${error instanceof Error ? error.message : String(error)}`, ) } @@ -104,21 +116,23 @@ export class ShadowCheckpointService extends CheckpointEventEmitter implements C // .git/info/exclude is local to the shadow git repo, so it's not // shared with the main repo - and won't conflict with user's // .gitignore. - await fs.mkdir(path.join(this.shadowGitDir, "info"), { recursive: true }) - const excludesPath = path.join(this.shadowGitDir, "info", "exclude") + await fs.mkdir(path.join(this.dotGitDir, "info"), { recursive: true }) + const excludesPath = path.join(this.dotGitDir, "info", "exclude") await fs.writeFile(excludesPath, [...GIT_EXCLUDES, ...lfsPatterns].join("\n")) - await this.stageAll() - const { commit } = await this.git.commit("initial commit", { "--allow-empty": null }) + await this.stageAll(git) + const { commit } = await git.commit("initial commit", { "--allow-empty": null }) this.baseHash = commit - this.log(`[CheckpointService#initShadowGit] base commit is ${commit}`) - created = true } const duration = Date.now() - startTime - this.log(`[CheckpointService#initShadowGit] initialized shadow git in ${duration}ms`) + this.log( + `[${this.constructor.name}#initShadowGit] initialized shadow repo with base commit ${this.baseHash} in ${duration}ms`, + ) - this.isInitialized = true + this.git = git + + await onInit?.() this.emit("initialize", { type: "initialize", @@ -127,16 +141,19 @@ export class ShadowCheckpointService extends CheckpointEventEmitter implements C created, duration, }) + + return { created, duration } } - private async stageAll() { + private async stageAll(git: SimpleGit) { + // await writeExcludesFile(gitPath, await getLfsPatterns(this.cwd)). await this.renameNestedGitRepos(true) try { - await this.git.add(".") + await git.add(".") } catch (error) { this.log( - `[CheckpointService#stageAll] failed to add files to git: ${error instanceof Error ? error.message : String(error)}`, + `[${this.constructor.name}#stageAll] failed to add files to git: ${error instanceof Error ? error.message : String(error)}`, ) } finally { await this.renameNestedGitRepos(false) @@ -172,23 +189,23 @@ export class ShadowCheckpointService extends CheckpointEventEmitter implements C try { await fs.rename(fullPath, newPath) this.log( - `[CheckpointService#renameNestedGitRepos] ${disable ? "disabled" : "enabled"} nested git repo ${gitPath}`, + `[${this.constructor.name}#renameNestedGitRepos] ${disable ? "disabled" : "enabled"} nested git repo ${gitPath}`, ) } catch (error) { this.log( - `[CheckpointService#renameNestedGitRepos] failed to ${disable ? "disable" : "enable"} nested git repo ${gitPath}: ${error instanceof Error ? error.message : String(error)}`, + `[${this.constructor.name}#renameNestedGitRepos] failed to ${disable ? "disable" : "enable"} nested git repo ${gitPath}: ${error instanceof Error ? error.message : String(error)}`, ) } } } - public async getShadowGitConfigWorktree() { + private async getShadowGitConfigWorktree(git: SimpleGit) { if (!this.shadowGitConfigWorktree) { try { - this.shadowGitConfigWorktree = (await this.git.getConfig("core.worktree")).value || undefined + this.shadowGitConfigWorktree = (await git.getConfig("core.worktree")).value || undefined } catch (error) { this.log( - `[CheckpointService#getShadowGitConfigWorktree] failed to get core.worktree: ${error instanceof Error ? error.message : String(error)}`, + `[${this.constructor.name}#getShadowGitConfigWorktree] failed to get core.worktree: ${error instanceof Error ? error.message : String(error)}`, ) } } @@ -196,27 +213,39 @@ export class ShadowCheckpointService extends CheckpointEventEmitter implements C return this.shadowGitConfigWorktree } - public async saveCheckpoint(message: string) { + public async saveCheckpoint(message: string): Promise { try { - this.log("[CheckpointService#saveCheckpoint] starting checkpoint save") + this.log(`[${this.constructor.name}#saveCheckpoint] starting checkpoint save`) - if (!this.isInitialized) { + if (!this.git) { throw new Error("Shadow git repo not initialized") } const startTime = Date.now() - await this.stageAll() + await this.stageAll(this.git) const result = await this.git.commit(message) const isFirst = this._checkpoints.length === 0 const fromHash = this._checkpoints[this._checkpoints.length - 1] ?? this.baseHash! - const toHash = result.commit ?? fromHash + const toHash = result.commit || fromHash this._checkpoints.push(toHash) const duration = Date.now() - startTime - this.emit("checkpoint", { type: "checkpoint", isFirst, fromHash, toHash, duration }) - return result.commit ? result : undefined + + if (isFirst || result.commit) { + this.emit("checkpoint", { type: "checkpoint", isFirst, fromHash, toHash, duration }) + } + + if (result.commit) { + this.log( + `[${this.constructor.name}#saveCheckpoint] checkpoint saved in ${duration}ms -> ${result.commit}`, + ) + return result + } else { + this.log(`[${this.constructor.name}#saveCheckpoint] found no changes to commit in ${duration}ms`) + return undefined + } } catch (e) { const error = e instanceof Error ? e : new Error(String(e)) - this.log(`[CheckpointService#saveCheckpoint] failed to create checkpoint: ${error.message}`) + this.log(`[${this.constructor.name}#saveCheckpoint] failed to create checkpoint: ${error.message}`) this.emit("error", { type: "error", error }) throw error } @@ -224,7 +253,9 @@ export class ShadowCheckpointService extends CheckpointEventEmitter implements C public async restoreCheckpoint(commitHash: string) { try { - if (!this.isInitialized) { + this.log(`[${this.constructor.name}#restoreCheckpoint] starting checkpoint restore`) + + if (!this.git) { throw new Error("Shadow git repo not initialized") } @@ -241,17 +272,17 @@ export class ShadowCheckpointService extends CheckpointEventEmitter implements C const duration = Date.now() - start this.emit("restore", { type: "restore", commitHash, duration }) - this.log(`[CheckpointService#restoreCheckpoint] restored checkpoint ${commitHash} in ${duration}ms`) + this.log(`[${this.constructor.name}#restoreCheckpoint] restored checkpoint ${commitHash} in ${duration}ms`) } catch (e) { const error = e instanceof Error ? e : new Error(String(e)) - this.log(`[CheckpointService#restoreCheckpoint] failed to restore checkpoint: ${error.message}`) + this.log(`[${this.constructor.name}#restoreCheckpoint] failed to restore checkpoint: ${error.message}`) this.emit("error", { type: "error", error }) throw error } } - public async getDiff({ from, to }: { from?: string; to?: string }) { - if (!this.isInitialized) { + public async getDiff({ from, to }: { from?: string; to?: string }): Promise { + if (!this.git) { throw new Error("Shadow git repo not initialized") } @@ -262,11 +293,12 @@ export class ShadowCheckpointService extends CheckpointEventEmitter implements C } // Stage all changes so that untracked files appear in diff summary. - await this.stageAll() + await this.stageAll(this.git) + this.log(`[${this.constructor.name}#getDiff] diffing ${to ? `${from}..${to}` : `${from}..HEAD`}`) const { files } = to ? await this.git.diffSummary([`${from}..${to}`]) : await this.git.diffSummary([from]) - const cwdPath = (await this.getShadowGitConfigWorktree()) || this.workspaceDir || "" + const cwdPath = (await this.getShadowGitConfigWorktree(this.git)) || this.workspaceDir || "" for (const file of files) { const relPath = file.file @@ -283,30 +315,23 @@ export class ShadowCheckpointService extends CheckpointEventEmitter implements C return result } - public static async create({ taskId, shadowDir, workspaceDir, log = console.log }: ShadowCheckpointServiceOptions) { - try { - await simpleGit().version() - } catch (error) { - log("[CheckpointService#create] git is not installed") - throw new Error("Git must be installed to use checkpoints.") - } + /** + * EventEmitter + */ - const homedir = os.homedir() - const desktopPath = path.join(homedir, "Desktop") - const documentsPath = path.join(homedir, "Documents") - const downloadsPath = path.join(homedir, "Downloads") - const protectedPaths = [homedir, desktopPath, documentsPath, downloadsPath] + override emit(event: K, data: CheckpointEventMap[K]) { + return super.emit(event, data) + } - if (protectedPaths.includes(workspaceDir)) { - throw new Error(`Cannot use checkpoints in ${workspaceDir}`) - } + override on(event: K, listener: (data: CheckpointEventMap[K]) => void) { + return super.on(event, listener) + } - const checkpointsDir = path.join(shadowDir, "tasks", taskId, "checkpoints") - await fs.mkdir(checkpointsDir, { recursive: true }) - const gitDir = path.join(checkpointsDir, ".git") - const git = simpleGit(path.dirname(gitDir)) + override off(event: K, listener: (data: CheckpointEventMap[K]) => void) { + return super.off(event, listener) + } - log(`[CheckpointService#create] taskId = ${taskId}, workspaceDir = ${workspaceDir}, shadowDir = ${shadowDir}`) - return new ShadowCheckpointService(taskId, git, shadowDir, workspaceDir, log) + override once(event: K, listener: (data: CheckpointEventMap[K]) => void) { + return super.once(event, listener) } } diff --git a/src/services/checkpoints/__tests__/ShadowCheckpointService.test.ts b/src/services/checkpoints/__tests__/ShadowCheckpointService.test.ts index 138e1130e8..bafcb3eeb2 100644 --- a/src/services/checkpoints/__tests__/ShadowCheckpointService.test.ts +++ b/src/services/checkpoints/__tests__/ShadowCheckpointService.test.ts @@ -7,65 +7,69 @@ import { EventEmitter } from "events" import { simpleGit, SimpleGit } from "simple-git" -import { ShadowCheckpointService } from "../ShadowCheckpointService" +import { RepoPerTaskCheckpointService } from "../RepoPerTaskCheckpointService" +import { RepoPerWorkspaceCheckpointService } from "../RepoPerWorkspaceCheckpointService" jest.mock("globby", () => ({ globby: jest.fn().mockResolvedValue([]), })) -const tmpDir = path.join(os.tmpdir(), "test-ShadowCheckpointService") +const tmpDir = path.join(os.tmpdir(), "CheckpointService") -describe("ShadowCheckpointService", () => { +const initRepo = async ({ + workspaceDir, + userName = "Roo Code", + userEmail = "support@roocode.com", + testFileName = "test.txt", + textFileContent = "Hello, world!", +}: { + workspaceDir: string + userName?: string + userEmail?: string + testFileName?: string + textFileContent?: string +}) => { + // Create a temporary directory for testing. + await fs.mkdir(workspaceDir, { recursive: true }) + + // Initialize git repo. + const git = simpleGit(workspaceDir) + await git.init() + await git.addConfig("user.name", userName) + await git.addConfig("user.email", userEmail) + + // Create test file. + const testFile = path.join(workspaceDir, testFileName) + await fs.writeFile(testFile, textFileContent) + + // Create initial commit. + await git.add(".") + await git.commit("Initial commit")! + + return { git, testFile } +} + +describe.each([ + [RepoPerTaskCheckpointService, "RepoPerTaskCheckpointService"], + [RepoPerWorkspaceCheckpointService, "RepoPerWorkspaceCheckpointService"], +])("CheckpointService", (klass, prefix) => { const taskId = "test-task" let workspaceGit: SimpleGit let testFile: string - let service: ShadowCheckpointService - - const initRepo = async ({ - workspaceDir, - userName = "Roo Code", - userEmail = "support@roocode.com", - testFileName = "test.txt", - textFileContent = "Hello, world!", - }: { - workspaceDir: string - userName?: string - userEmail?: string - testFileName?: string - textFileContent?: string - }) => { - // Create a temporary directory for testing. - await fs.mkdir(workspaceDir, { recursive: true }) - - // Initialize git repo. - const git = simpleGit(workspaceDir) - await git.init() - await git.addConfig("user.name", userName) - await git.addConfig("user.email", userEmail) - - // Create test file. - const testFile = path.join(workspaceDir, testFileName) - await fs.writeFile(testFile, textFileContent) - - // Create initial commit. - await git.add(".") - await git.commit("Initial commit")! - - return { git, testFile } - } + let service: RepoPerTaskCheckpointService | RepoPerWorkspaceCheckpointService beforeEach(async () => { jest.mocked(require("globby").globby).mockClear().mockResolvedValue([]) - const shadowDir = path.join(tmpDir, `shadow-${Date.now()}`) + const shadowDir = path.join(tmpDir, `${prefix}-${Date.now()}`) const workspaceDir = path.join(tmpDir, `workspace-${Date.now()}`) const repo = await initRepo({ workspaceDir }) - testFile = repo.testFile workspaceGit = repo.git + testFile = repo.testFile - service = await ShadowCheckpointService.create({ taskId, shadowDir, workspaceDir, log: () => {} }) + service = await klass.create({ taskId, shadowDir, workspaceDir, log: () => {} }) await service.initShadowGit() }) @@ -77,14 +81,14 @@ describe("ShadowCheckpointService", () => { await fs.rm(tmpDir, { recursive: true, force: true }) }) - describe("getDiff", () => { + describe(`${klass.name}#getDiff`, () => { it("returns the correct diff between commits", async () => { await fs.writeFile(testFile, "Ahoy, world!") - const commit1 = await service.saveCheckpoint("First checkpoint") + const commit1 = await service.saveCheckpoint("Ahoy, world!") expect(commit1?.commit).toBeTruthy() await fs.writeFile(testFile, "Goodbye, world!") - const commit2 = await service.saveCheckpoint("Second checkpoint") + const commit2 = await service.saveCheckpoint("Goodbye, world!") expect(commit2?.commit).toBeTruthy() const diff1 = await service.getDiff({ to: commit1!.commit }) @@ -94,7 +98,7 @@ describe("ShadowCheckpointService", () => { expect(diff1[0].content.before).toBe("Hello, world!") expect(diff1[0].content.after).toBe("Ahoy, world!") - const diff2 = await service.getDiff({ to: commit2!.commit }) + const diff2 = await service.getDiff({ from: service.baseHash, to: commit2!.commit }) expect(diff2).toHaveLength(1) expect(diff2[0].paths.relative).toBe("test.txt") expect(diff2[0].paths.absolute).toBe(testFile) @@ -140,7 +144,7 @@ describe("ShadowCheckpointService", () => { }) }) - describe("saveCheckpoint", () => { + describe(`${klass.name}#saveCheckpoint`, () => { it("creates a checkpoint if there are pending changes", async () => { await fs.writeFile(testFile, "Ahoy, world!") const commit1 = await service.saveCheckpoint("First checkpoint") @@ -296,9 +300,9 @@ describe("ShadowCheckpointService", () => { }) }) - describe("create", () => { + describe(`${klass.name}#create`, () => { it("initializes a git repository if one does not already exist", async () => { - const shadowDir = path.join(tmpDir, `shadow2-${Date.now()}`) + const shadowDir = path.join(tmpDir, `${prefix}2-${Date.now()}`) const workspaceDir = path.join(tmpDir, `workspace2-${Date.now()}`) await fs.mkdir(workspaceDir) @@ -307,10 +311,11 @@ describe("ShadowCheckpointService", () => { expect(await fs.readFile(newTestFile, "utf-8")).toBe("Hello, world!") // Ensure the git repository was initialized. - const gitDir = path.join(shadowDir, "tasks", taskId, "checkpoints", ".git") - await expect(fs.stat(gitDir)).rejects.toThrow() - const newService = await ShadowCheckpointService.create({ taskId, shadowDir, workspaceDir, log: () => {} }) - await newService.initShadowGit() + const newService = await klass.create({ taskId, shadowDir, workspaceDir, log: () => {} }) + const { created } = await newService.initShadowGit() + expect(created).toBeTruthy() + + const gitDir = path.join(newService.checkpointsDir, ".git") expect(await fs.stat(gitDir)).toBeTruthy() // Save a new checkpoint: Ahoy, world! @@ -327,15 +332,15 @@ describe("ShadowCheckpointService", () => { await newService.restoreCheckpoint(commit1!.commit) expect(await fs.readFile(newTestFile, "utf-8")).toBe("Ahoy, world!") - await fs.rm(newService.shadowDir, { recursive: true, force: true }) + await fs.rm(newService.checkpointsDir, { recursive: true, force: true }) await fs.rm(newService.workspaceDir, { recursive: true, force: true }) }) }) - describe("events", () => { + describe(`${klass.name}#events`, () => { it("emits initialize event when service is created", async () => { - const shadowDir = path.join(tmpDir, `shadow-event-test-${Date.now()}`) - const workspaceDir = path.join(tmpDir, `workspace-event-test-${Date.now()}`) + const shadowDir = path.join(tmpDir, `${prefix}3-${Date.now()}`) + const workspaceDir = path.join(tmpDir, `workspace3-${Date.now()}`) await fs.mkdir(workspaceDir, { recursive: true }) const newTestFile = path.join(workspaceDir, "test.txt") @@ -345,7 +350,7 @@ describe("ShadowCheckpointService", () => { const emitSpy = jest.spyOn(EventEmitter.prototype, "emit") // Create the service - this will trigger the initialize event. - const newService = await ShadowCheckpointService.create({ taskId, shadowDir, workspaceDir, log: () => {} }) + const newService = await klass.create({ taskId, shadowDir, workspaceDir, log: () => {} }) await newService.initShadowGit() // Find the initialize event in the emit calls. diff --git a/src/services/checkpoints/index.ts b/src/services/checkpoints/index.ts new file mode 100644 index 0000000000..9794b34d4c --- /dev/null +++ b/src/services/checkpoints/index.ts @@ -0,0 +1,4 @@ +export type { CheckpointServiceOptions } from "./types" + +export { RepoPerTaskCheckpointService } from "./RepoPerTaskCheckpointService" +export { RepoPerWorkspaceCheckpointService } from "./RepoPerWorkspaceCheckpointService" diff --git a/src/services/checkpoints/types.ts b/src/services/checkpoints/types.ts index 3ae9457fe8..81611e81ec 100644 --- a/src/services/checkpoints/types.ts +++ b/src/services/checkpoints/types.ts @@ -1,5 +1,4 @@ -import EventEmitter from "events" -import { CommitResult } from "simple-git" +import { CommitResult, SimpleGit } from "simple-git" export type CheckpointResult = Partial & Pick @@ -14,25 +13,14 @@ export type CheckpointDiff = { } } -export interface CheckpointService { - saveCheckpoint(message: string): Promise - restoreCheckpoint(commit: string): Promise - getDiff(range: { from?: string; to?: string }): Promise - workspaceDir: string - baseHash?: string - version: number -} - export interface CheckpointServiceOptions { taskId: string workspaceDir: string + shadowDir: string // globalStorageUri.fsPath + log?: (message: string) => void } -/** - * EventEmitter - */ - export interface CheckpointEventMap { initialize: { type: "initialize"; workspaceDir: string; baseHash: string; created: boolean; duration: number } checkpoint: { @@ -45,21 +33,3 @@ export interface CheckpointEventMap { restore: { type: "restore"; commitHash: string; duration: number } error: { type: "error"; error: Error } } - -export class CheckpointEventEmitter extends EventEmitter { - override emit(event: K, data: CheckpointEventMap[K]): boolean { - return super.emit(event, data) - } - - override on(event: K, listener: (data: CheckpointEventMap[K]) => void): this { - return super.on(event, listener) - } - - override off(event: K, listener: (data: CheckpointEventMap[K]) => void): this { - return super.off(event, listener) - } - - override once(event: K, listener: (data: CheckpointEventMap[K]) => void): this { - return super.once(event, listener) - } -} diff --git a/src/shared/WebviewMessage.ts b/src/shared/WebviewMessage.ts index 3b1750bcb6..141915fdb8 100644 --- a/src/shared/WebviewMessage.ts +++ b/src/shared/WebviewMessage.ts @@ -122,6 +122,7 @@ export interface WebviewMessage { export const checkoutDiffPayloadSchema = z.object({ ts: z.number(), + previousCommitHash: z.string().optional(), commitHash: z.string(), mode: z.enum(["full", "checkpoint"]), }) diff --git a/webview-ui/src/components/chat/checkpoints/CheckpointMenu.tsx b/webview-ui/src/components/chat/checkpoints/CheckpointMenu.tsx index aa52c2db88..1910fc0bb6 100644 --- a/webview-ui/src/components/chat/checkpoints/CheckpointMenu.tsx +++ b/webview-ui/src/components/chat/checkpoints/CheckpointMenu.tsx @@ -20,13 +20,17 @@ export const CheckpointMenu = ({ ts, commitHash, currentHash, checkpoint }: Chec const isCurrent = currentHash === commitHash const isFirst = checkpoint.isFirst - const isDiffAvailable = !isFirst const isRestoreAvailable = !isFirst || !isCurrent + const previousCommitHash = checkpoint?.from + const onCheckpointDiff = useCallback(() => { - vscode.postMessage({ type: "checkpointDiff", payload: { ts, commitHash, mode: "checkpoint" } }) - }, [ts, commitHash]) + vscode.postMessage({ + type: "checkpointDiff", + payload: { ts, previousCommitHash, commitHash, mode: "checkpoint" }, + }) + }, [ts, previousCommitHash, commitHash]) const onPreview = useCallback(() => { vscode.postMessage({ type: "checkpointRestore", payload: { ts, commitHash, mode: "preview" } }) diff --git a/webview-ui/src/components/history/HistoryView.tsx b/webview-ui/src/components/history/HistoryView.tsx index d50a569c8d..c621da5a03 100644 --- a/webview-ui/src/components/history/HistoryView.tsx +++ b/webview-ui/src/components/history/HistoryView.tsx @@ -299,7 +299,12 @@ const HistoryView = ({ onDone }: HistoryViewProps) => { {formatLargeNumber(item.tokensOut || 0)} - {!item.totalCost && } + {!item.totalCost && ( +
+ + +
+ )} {!!item.cacheWrites && ( From d7fb9980e6f99cee8594a2738ab9f1c909f7a06d Mon Sep 17 00:00:00 2001 From: cte Date: Tue, 4 Mar 2025 01:39:20 -0800 Subject: [PATCH 2/2] Revert logging in ClineProvider --- src/core/webview/ClineProvider.ts | 4 ---- 1 file changed, 4 deletions(-) diff --git a/src/core/webview/ClineProvider.ts b/src/core/webview/ClineProvider.ts index d318989dfa..0a2e385b6a 100644 --- a/src/core/webview/ClineProvider.ts +++ b/src/core/webview/ClineProvider.ts @@ -887,8 +887,6 @@ export class ClineProvider implements vscode.WebviewViewProvider { if (result.success) { await this.cline?.checkpointDiff(result.data) - } else { - vscode.window.showErrorMessage("Invalid checkpoint diff payload.") } break @@ -909,8 +907,6 @@ export class ClineProvider implements vscode.WebviewViewProvider { } catch (error) { vscode.window.showErrorMessage("Failed to restore checkpoint.") } - } else { - vscode.window.showErrorMessage("Invalid checkpoint restore payload.") } break