diff --git a/src/core/Cline.ts b/src/core/Cline.ts index c66dec2c86..b9b243dbc1 100644 --- a/src/core/Cline.ts +++ b/src/core/Cline.ts @@ -10,10 +10,10 @@ 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 { ApiHandler, buildApiHandler } from "../api" import { ApiStream } from "../api/transform/stream" import { DIFF_VIEW_URI_SCHEME, DiffViewProvider } from "../integrations/editor/DiffViewProvider" -import { CheckpointService, CheckpointServiceFactory } from "../services/checkpoints" +import { ShadowCheckpointService } from "../services/checkpoints/ShadowCheckpointService" import { findToolName, formatContentBlockToMarkdown } from "../integrations/misc/export-markdown" import { extractTextFromFile, @@ -116,7 +116,7 @@ export class Cline { // checkpoints enableCheckpoints: boolean = false - private checkpointService?: CheckpointService + private checkpointService?: ShadowCheckpointService // streaming isWaitingForFirstChunk = false @@ -747,8 +747,11 @@ export class Cline { } private async initiateTaskLoop(userContent: UserContent): Promise { + this.initializeCheckpoints() + let nextUserContent = userContent let includeFileDetails = true + while (!this.abort) { const didEndLoop = await this.recursivelyMakeClineRequests(nextUserContent, includeFileDetails) includeFileDetails = false // we only need file details the first time @@ -2773,7 +2776,7 @@ export class Cline { } if (isCheckpointPossible) { - await this.checkpointSave({ isFirst: false }) + this.checkpointSave() } /* @@ -2839,13 +2842,6 @@ export class Cline { // get previous api req's index to check token usage and determine if we need to truncate conversation history const previousApiReqIndex = findLastIndex(this.clineMessages, (m) => m.say === "api_req_started") - // Save checkpoint if this is the first API request. - const isFirstRequest = this.clineMessages.filter((m) => m.say === "api_req_started").length === 0 - - if (isFirstRequest) { - 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 // for the best UX we show a placeholder api_req_started message with a loading spinner as this happens await this.say( @@ -3356,37 +3352,88 @@ export class Cline { // Checkpoints - private async getCheckpointService() { + private async initializeCheckpoints() { if (!this.enableCheckpoints) { - throw new Error("Checkpoints are disabled") + return } - if (!this.checkpointService) { + const log = (message: string) => { + console.log(message) + + try { + this.providerRef.deref()?.log(message) + } catch (err) { + // NO-OP + } + } + + try { + if (this.checkpointService) { + log("[Cline#initializeCheckpoints] checkpointService already initialized") + return + } + 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") + log("[Cline#initializeCheckpoints] workspace folder not found, disabling checkpoints") + this.enableCheckpoints = false + return } + const shadowDir = this.providerRef.deref()?.context.globalStorageUri.fsPath + if (!shadowDir) { - this.providerRef.deref()?.log("[getCheckpointService] shadowDir not found") - throw new Error("Global storage directory not found") + log("[Cline#initializeCheckpoints] shadowDir not found, disabling checkpoints") + this.enableCheckpoints = false + return } - this.checkpointService = await CheckpointServiceFactory.create({ - strategy: "shadow", - options: { - taskId: this.taskId, - workspaceDir, - shadowDir, - log: (message) => this.providerRef.deref()?.log(message), - }, - }) - } + const service = await ShadowCheckpointService.create({ taskId: this.taskId, workspaceDir, shadowDir, log }) - return this.checkpointService + if (!service) { + log("[Cline#initializeCheckpoints] failed to create checkpoint service, disabling checkpoints") + this.enableCheckpoints = false + return + } + + 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`, + ) + } + + this.checkpointService = service + this.checkpointSave() + } catch (err) { + log("[Cline#initializeCheckpoints] caught error in on('initialize'), disabling checkpoints") + this.enableCheckpoints = false + } + }) + + 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), + ) + } catch (err) { + log("[Cline#initializeCheckpoints] caught error in on('checkpoint'), disabling checkpoints") + this.enableCheckpoints = false + } + }) + + service.initShadowGit() + } catch (err) { + log("[Cline#initializeCheckpoints] caught error in initializeCheckpoints(), disabling checkpoints") + this.enableCheckpoints = false + } } public async checkpointDiff({ @@ -3398,7 +3445,7 @@ export class Cline { commitHash: string mode: "full" | "checkpoint" }) { - if (!this.enableCheckpoints) { + if (!this.checkpointService || !this.enableCheckpoints) { return } @@ -3414,8 +3461,7 @@ export class Cline { } try { - const service = await this.getCheckpointService() - const changes = await service.getDiff({ from: previousCommitHash, to: commitHash }) + const changes = await this.checkpointService.getDiff({ from: previousCommitHash, to: commitHash }) if (!changes?.length) { vscode.window.showInformationMessage("No changes found.") @@ -3441,30 +3487,13 @@ export class Cline { } } - public async checkpointSave({ isFirst }: { isFirst: boolean }) { - if (!this.enableCheckpoints) { + public checkpointSave() { + if (!this.checkpointService || !this.enableCheckpoints) { return } - try { - const service = await this.getCheckpointService() - const strategy = service.strategy - const version = service.version - - const commit = await service.saveCheckpoint(`Task: ${this.taskId}, Time: ${Date.now()}`) - const fromHash = service.baseHash - const toHash = isFirst ? commit?.commit || fromHash : commit?.commit - - if (toHash) { - await this.providerRef.deref()?.postMessageToWebview({ type: "currentCheckpointUpdated", text: toHash }) - - 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") - this.enableCheckpoints = false - } + // Start the checkpoint process in the background. + this.checkpointService.saveCheckpoint(`Task: ${this.taskId}, Time: ${Date.now()}`) } public async checkpointRestore({ @@ -3476,7 +3505,7 @@ export class Cline { commitHash: string mode: "preview" | "restore" }) { - if (!this.enableCheckpoints) { + if (!this.checkpointService || !this.enableCheckpoints) { return } @@ -3487,8 +3516,7 @@ export class Cline { } try { - const service = await this.getCheckpointService() - await service.restoreCheckpoint(commitHash) + await this.checkpointService.restoreCheckpoint(commitHash) await this.providerRef.deref()?.postMessageToWebview({ type: "currentCheckpointUpdated", text: commitHash }) diff --git a/src/services/checkpoints/CheckpointServiceFactory.ts b/src/services/checkpoints/CheckpointServiceFactory.ts deleted file mode 100644 index ff39a1fe7e..0000000000 --- a/src/services/checkpoints/CheckpointServiceFactory.ts +++ /dev/null @@ -1,29 +0,0 @@ -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 { strategy: "local" } - ? LocalCheckpointService - : T extends { strategy: "shadow" } - ? ShadowCheckpointService - : never - -export class CheckpointServiceFactory { - public static create(options: T): CheckpointServiceType { - switch (options.strategy) { - case "local": - return LocalCheckpointService.create(options.options) as any - case "shadow": - return ShadowCheckpointService.create(options.options) as any - } - } -} diff --git a/src/services/checkpoints/LocalCheckpointService.ts b/src/services/checkpoints/LocalCheckpointService.ts deleted file mode 100644 index ce5c6bd6ea..0000000000 --- a/src/services/checkpoints/LocalCheckpointService.ts +++ /dev/null @@ -1,440 +0,0 @@ -import fs from "fs/promises" -import { existsSync } from "fs" -import path from "path" - -import simpleGit, { SimpleGit, CleanOptions } from "simple-git" - -import { CheckpointStrategy, CheckpointService, CheckpointServiceOptions } from "./types" - -export interface LocalCheckpointServiceOptions extends CheckpointServiceOptions {} - -/** - * The CheckpointService provides a mechanism for storing a snapshot of the - * current VSCode workspace each time a Roo Code tool is executed. It uses Git - * under the hood. - * - * HOW IT WORKS - * - * Two branches are used: - * - A main branch for normal operation (the branch you are currently on). - * - A hidden branch for storing checkpoints. - * - * Saving a checkpoint: - * - A temporary branch is created to store the current state. - * - All changes (including untracked files) are staged and committed on the temp branch. - * - The hidden branch is reset to match main. - * - The temporary branch commit is cherry-picked onto the hidden branch. - * - The workspace is restored to its original state and the temp branch is deleted. - * - * Restoring a checkpoint: - * - The workspace is restored to the state of the specified checkpoint using - * `git restore` and `git clean`. - * - * This approach allows for: - * - Non-destructive version control (main branch remains untouched). - * - Preservation of the full history of checkpoints. - * - Safe restoration to any previous checkpoint. - * - Atomic checkpoint operations with proper error recovery. - * - * NOTES - * - * - Git must be installed. - * - If the current working directory is not a Git repository, we will - * initialize a new one with a .gitkeep file. - * - If you manually edit files and then restore a checkpoint, the changes - * will be lost. Addressing this adds some complexity to the implementation - * and it's not clear whether it's worth it. - */ - -export class LocalCheckpointService implements CheckpointService { - private static readonly USER_NAME = "Roo Code" - private static readonly USER_EMAIL = "support@roocode.com" - private static readonly CHECKPOINT_BRANCH = "roo-code-checkpoints" - private static readonly STASH_BRANCH = "roo-code-stash" - - public readonly strategy: CheckpointStrategy = "local" - public readonly version = 1 - - public get baseHash() { - return this._baseHash - } - - constructor( - public readonly taskId: string, - public readonly git: SimpleGit, - public readonly workspaceDir: string, - private readonly mainBranch: string, - private _baseHash: string, - private readonly hiddenBranch: string, - private readonly log: (message: string) => void, - ) {} - - private async ensureBranch(expectedBranch: string) { - const branch = await this.git.revparse(["--abbrev-ref", "HEAD"]) - - if (branch.trim() !== expectedBranch) { - throw new Error(`Git branch mismatch: expected '${expectedBranch}' but found '${branch}'`) - } - } - - public async getDiff({ from, to }: { from?: string; to?: string }) { - const result = [] - - if (!from) { - from = this.baseHash - } - - const { files } = await this.git.diffSummary([`${from}..${to}`]) - - for (const file of files.filter((f) => !f.binary)) { - const relPath = file.file - const absPath = path.join(this.workspaceDir, relPath) - const before = await this.git.show([`${from}:${relPath}`]).catch(() => "") - - const after = to - ? await this.git.show([`${to}:${relPath}`]).catch(() => "") - : await fs.readFile(absPath, "utf8").catch(() => "") - - result.push({ - paths: { relative: relPath, absolute: absPath }, - content: { before, after }, - }) - } - - return result - } - - private async restoreMain({ - branch, - stashSha, - force = false, - }: { - branch: string - stashSha: string - force?: boolean - }) { - let currentBranch = await this.git.revparse(["--abbrev-ref", "HEAD"]) - - if (currentBranch !== this.mainBranch) { - if (force) { - try { - await this.git.checkout(["-f", this.mainBranch]) - } catch (err) { - this.log( - `[restoreMain] failed to force checkout ${this.mainBranch}: ${err instanceof Error ? err.message : String(err)}`, - ) - } - } else { - try { - await this.git.checkout(this.mainBranch) - } catch (err) { - this.log( - `[restoreMain] failed to checkout ${this.mainBranch}: ${err instanceof Error ? err.message : String(err)}`, - ) - - // Escalate to a forced checkout if we can't checkout the - // main branch under normal circumstances. - currentBranch = await this.git.revparse(["--abbrev-ref", "HEAD"]) - - if (currentBranch !== this.mainBranch) { - await this.git.checkout(["-f", this.mainBranch]).catch(() => {}) - } - } - } - } - - currentBranch = await this.git.revparse(["--abbrev-ref", "HEAD"]) - - if (currentBranch !== this.mainBranch) { - throw new Error(`Unable to restore ${this.mainBranch}`) - } - - if (stashSha) { - this.log(`[restoreMain] applying stash ${stashSha}`) - - try { - await this.git.raw(["stash", "apply", "--index", stashSha]) - } catch (err) { - this.log(`[restoreMain] Failed to apply stash: ${err instanceof Error ? err.message : String(err)}`) - } - } - - this.log(`[restoreMain] restoring from ${branch} branch`) - - try { - await this.git.raw(["restore", "--source", branch, "--worktree", "--", "."]) - } catch (err) { - this.log(`[restoreMain] Failed to restore branch: ${err instanceof Error ? err.message : String(err)}`) - } - } - - public async saveCheckpoint(message: string) { - const startTime = Date.now() - - await this.ensureBranch(this.mainBranch) - - const stashSha = (await this.git.raw(["stash", "create"])).trim() - const latestSha = await this.git.revparse([this.hiddenBranch]) - - /** - * PHASE: Create stash - * Mutations: - * - Create branch - * - Change branch - */ - const stashBranch = `${LocalCheckpointService.STASH_BRANCH}-${Date.now()}` - await this.git.checkout(["-b", stashBranch]) - this.log(`[saveCheckpoint] created and checked out ${stashBranch}`) - - /** - * Phase: Stage stash - * Mutations: None - * Recovery: - * - UNDO: Create branch - * - UNDO: Change branch - */ - try { - await this.git.add(["-A"]) - } catch (err) { - this.log( - `[saveCheckpoint] failed in stage stash phase: ${err instanceof Error ? err.message : String(err)}`, - ) - await this.restoreMain({ branch: stashBranch, stashSha, force: true }) - await this.git.branch(["-D", stashBranch]).catch(() => {}) - throw err - } - - /** - * Phase: Commit stash - * Mutations: - * - Commit stash - * - Change branch - * Recovery: - * - UNDO: Create branch - * - UNDO: Change branch - */ - let stashCommit - - try { - stashCommit = await this.git.commit(message, undefined, { "--no-verify": null }) - this.log(`[saveCheckpoint] stashCommit: ${message} -> ${JSON.stringify(stashCommit)}`) - } catch (err) { - this.log( - `[saveCheckpoint] failed in stash commit phase: ${err instanceof Error ? err.message : String(err)}`, - ) - await this.restoreMain({ branch: stashBranch, stashSha, force: true }) - await this.git.branch(["-D", stashBranch]).catch(() => {}) - throw err - } - - if (!stashCommit) { - this.log("[saveCheckpoint] no stash commit") - await this.restoreMain({ branch: stashBranch, stashSha }) - await this.git.branch(["-D", stashBranch]) - return undefined - } - - /** - * PHASE: Diff - * Mutations: - * - Checkout hidden branch - * Recovery: - * - UNDO: Create branch - * - UNDO: Change branch - * - UNDO: Commit stash - */ - let diff - - try { - diff = await this.git.diff([latestSha, stashBranch]) - } catch (err) { - this.log(`[saveCheckpoint] failed in diff phase: ${err instanceof Error ? err.message : String(err)}`) - await this.restoreMain({ branch: stashBranch, stashSha, force: true }) - await this.git.branch(["-D", stashBranch]).catch(() => {}) - throw err - } - - if (!diff) { - this.log("[saveCheckpoint] no diff") - await this.restoreMain({ branch: stashBranch, stashSha }) - await this.git.branch(["-D", stashBranch]) - return undefined - } - - /** - * PHASE: Reset - * Mutations: - * - Reset hidden branch - * Recovery: - * - UNDO: Create branch - * - UNDO: Change branch - * - UNDO: Commit stash - */ - try { - await this.git.checkout(this.hiddenBranch) - this.log(`[saveCheckpoint] checked out ${this.hiddenBranch}`) - await this.git.reset(["--hard", this.mainBranch]) - this.log(`[saveCheckpoint] reset ${this.hiddenBranch}`) - } catch (err) { - this.log(`[saveCheckpoint] failed in reset phase: ${err instanceof Error ? err.message : String(err)}`) - await this.restoreMain({ branch: stashBranch, stashSha, force: true }) - await this.git.branch(["-D", stashBranch]).catch(() => {}) - throw err - } - - /** - * PHASE: Cherry pick - * Mutations: - * - Hidden commit (NOTE: reset on hidden branch no longer needed in - * success scenario.) - * Recovery: - * - UNDO: Create branch - * - UNDO: Change branch - * - UNDO: Commit stash - * - UNDO: Reset hidden branch - */ - let commit = "" - - try { - try { - await this.git.raw(["cherry-pick", stashBranch]) - } catch (err) { - // Check if we're in the middle of a cherry-pick. - // If the cherry-pick resulted in an empty commit (e.g., only - // deletions) then complete it with --allow-empty. - // Otherwise, rethrow the error. - if (existsSync(path.join(this.workspaceDir, ".git/CHERRY_PICK_HEAD"))) { - await this.git.raw(["commit", "--allow-empty", "--no-edit"]) - } else { - throw err - } - } - - commit = await this.git.revparse(["HEAD"]) - this.log(`[saveCheckpoint] cherry-pick commit = ${commit}`) - } catch (err) { - this.log( - `[saveCheckpoint] failed in cherry pick phase: ${err instanceof Error ? err.message : String(err)}`, - ) - await this.git.reset(["--hard", latestSha]).catch(() => {}) - await this.restoreMain({ branch: stashBranch, stashSha, force: true }) - await this.git.branch(["-D", stashBranch]).catch(() => {}) - throw err - } - - await this.restoreMain({ branch: stashBranch, stashSha }) - await this.git.branch(["-D", stashBranch]) - - // We've gotten reports that checkpoints can be slow in some cases, so - // we'll log the duration of the checkpoint save. - const duration = Date.now() - startTime - this.log(`[saveCheckpoint] saved checkpoint ${commit} in ${duration}ms`) - - return { commit } - } - - public async restoreCheckpoint(commitHash: string) { - const startTime = Date.now() - await this.ensureBranch(this.mainBranch) - await this.git.clean([CleanOptions.FORCE, CleanOptions.RECURSIVE]) - await this.git.raw(["restore", "--source", commitHash, "--worktree", "--", "."]) - const duration = Date.now() - startTime - this.log(`[restoreCheckpoint] restored checkpoint ${commitHash} in ${duration}ms`) - } - - public static async create({ taskId, workspaceDir, log = console.log }: LocalCheckpointServiceOptions) { - const git = simpleGit(workspaceDir) - const version = await git.version() - - if (!version?.installed) { - throw new Error(`Git is not installed. Please install Git if you wish to use checkpoints.`) - } - - if (!workspaceDir || !existsSync(workspaceDir)) { - throw new Error(`Base directory is not set or does not exist.`) - } - - const { currentBranch, currentSha, hiddenBranch } = await LocalCheckpointService.initRepo(git, { - taskId, - workspaceDir, - log, - }) - - log( - `[create] taskId = ${taskId}, workspaceDir = ${workspaceDir}, currentBranch = ${currentBranch}, currentSha = ${currentSha}, hiddenBranch = ${hiddenBranch}`, - ) - - return new LocalCheckpointService(taskId, git, workspaceDir, currentBranch, currentSha, hiddenBranch, log) - } - - private static async initRepo( - git: SimpleGit, - { taskId, workspaceDir, log }: Required, - ) { - const isExistingRepo = existsSync(path.join(workspaceDir, ".git")) - - if (!isExistingRepo) { - await git.init() - log(`[initRepo] Initialized new Git repository at ${workspaceDir}`) - } - - const globalUserName = await git.getConfig("user.name", "global") - const localUserName = await git.getConfig("user.name", "local") - const userName = localUserName.value || globalUserName.value - - const globalUserEmail = await git.getConfig("user.email", "global") - const localUserEmail = await git.getConfig("user.email", "local") - const userEmail = localUserEmail.value || globalUserEmail.value - - // Prior versions of this service indiscriminately set the local user - // config, and it should not override the global config. To address - // this we remove the local user config if it matches the default - // user name and email and there's a global config. - if (globalUserName.value && localUserName.value === LocalCheckpointService.USER_NAME) { - await git.raw(["config", "--unset", "--local", "user.name"]) - } - - if (globalUserEmail.value && localUserEmail.value === LocalCheckpointService.USER_EMAIL) { - await git.raw(["config", "--unset", "--local", "user.email"]) - } - - // Only set user config if not already configured. - if (!userName) { - await git.addConfig("user.name", LocalCheckpointService.USER_NAME) - } - - if (!userEmail) { - await git.addConfig("user.email", LocalCheckpointService.USER_EMAIL) - } - - if (!isExistingRepo) { - // We need at least one file to commit, otherwise the initial - // commit will fail, unless we use the `--allow-empty` flag. - // However, using an empty commit causes problems when restoring - // the checkpoint (i.e. the `git restore` command doesn't work - // for empty commits). - await fs.writeFile(path.join(workspaceDir, ".gitkeep"), "") - await git.add(".gitkeep") - const commit = await git.commit("Initial commit") - - if (!commit.commit) { - throw new Error("Failed to create initial commit") - } - - log(`[initRepo] Initial commit: ${commit.commit}`) - } - - const currentBranch = await git.revparse(["--abbrev-ref", "HEAD"]) - const currentSha = await git.revparse(["HEAD"]) - - const hiddenBranch = `${LocalCheckpointService.CHECKPOINT_BRANCH}-${taskId}` - const branchSummary = await git.branch() - - if (!branchSummary.all.includes(hiddenBranch)) { - await git.checkoutBranch(hiddenBranch, currentBranch) - await git.checkout(currentBranch) - } - - return { currentBranch, currentSha, hiddenBranch } - } -} diff --git a/src/services/checkpoints/ShadowCheckpointService.ts b/src/services/checkpoints/ShadowCheckpointService.ts index 301602312e..70b3c68484 100644 --- a/src/services/checkpoints/ShadowCheckpointService.ts +++ b/src/services/checkpoints/ShadowCheckpointService.ts @@ -5,17 +5,18 @@ import { globby } from "globby" import simpleGit, { SimpleGit } from "simple-git" import { GIT_DISABLED_SUFFIX, GIT_EXCLUDES } from "./constants" -import { CheckpointStrategy, CheckpointService, CheckpointServiceOptions } from "./types" +import { CheckpointService, CheckpointServiceOptions, CheckpointEventEmitter } from "./types" export interface ShadowCheckpointServiceOptions extends CheckpointServiceOptions { shadowDir: string } -export class ShadowCheckpointService implements CheckpointService { - public readonly strategy: CheckpointStrategy = "shadow" +export class ShadowCheckpointService extends CheckpointEventEmitter implements CheckpointService { public readonly version = 1 + private _checkpoints: string[] = [] private _baseHash?: string + private _isInitialized = false public get baseHash() { return this._baseHash @@ -25,6 +26,18 @@ export class ShadowCheckpointService implements CheckpointService { this._baseHash = value } + public get isInitialized() { + return this._isInitialized + } + + private set isInitialized(value: boolean) { + this._isInitialized = value + } + + public get isDirty() { + return this._checkpoints.length > 0 + } + private readonly shadowGitDir: string private shadowGitConfigWorktree?: string @@ -35,18 +48,26 @@ export class ShadowCheckpointService implements CheckpointService { public readonly workspaceDir: string, private readonly log: (message: string) => void, ) { + super() this.shadowGitDir = path.join(this.shadowDir, "tasks", this.taskId, "checkpoints", ".git") } - private async initShadowGit() { + public async initShadowGit() { + if (this.isInitialized) { + return + } + const fileExistsAtPath = (path: string) => fs .access(path) .then(() => true) .catch(() => false) + let created = false + const startTime = Date.now() + if (await fileExistsAtPath(this.shadowGitDir)) { - this.log(`[initShadowGit] shadow git repo already exists at ${this.shadowGitDir}`) + this.log(`[CheckpointService#initShadowGit] shadow git repo already exists at ${this.shadowGitDir}`) const worktree = await this.getShadowGitConfigWorktree() if (worktree !== this.workspaceDir) { @@ -57,7 +78,7 @@ export class ShadowCheckpointService implements CheckpointService { this.baseHash = await this.git.revparse(["--abbrev-ref", "HEAD"]) } else { - this.log(`[initShadowGit] creating shadow git repo at ${this.workspaceDir}`) + this.log(`[CheckpointService#initShadowGit] creating shadow git repo at ${this.workspaceDir}`) await this.git.init() await this.git.addConfig("core.worktree", this.workspaceDir) // Sets the working tree to the current workspace. @@ -78,7 +99,7 @@ export class ShadowCheckpointService implements CheckpointService { } } catch (error) { this.log( - `[initShadowGit] failed to read .gitattributes: ${error instanceof Error ? error.message : String(error)}`, + `[CheckpointService#initShadowGit] failed to read .gitattributes: ${error instanceof Error ? error.message : String(error)}`, ) } @@ -93,8 +114,23 @@ export class ShadowCheckpointService implements CheckpointService { await this.stageAll() const { commit } = await this.git.commit("initial commit", { "--allow-empty": null }) this.baseHash = commit - this.log(`[initShadowGit] base commit is ${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.isInitialized = true + + this.emit("initialize", { + type: "initialize", + workspaceDir: this.workspaceDir, + baseHash: this.baseHash, + created, + duration, + }) } private async stageAll() { @@ -103,7 +139,9 @@ export class ShadowCheckpointService implements CheckpointService { try { await this.git.add(".") } catch (error) { - this.log(`[stageAll] failed to add files to git: ${error instanceof Error ? error.message : String(error)}`) + this.log( + `[CheckpointService#stageAll] failed to add files to git: ${error instanceof Error ? error.message : String(error)}`, + ) } finally { await this.renameNestedGitRepos(false) } @@ -137,10 +175,12 @@ export class ShadowCheckpointService implements CheckpointService { try { await fs.rename(fullPath, newPath) - this.log(`${disable ? "disabled" : "enabled"} nested git repo ${gitPath}`) + this.log( + `[CheckpointService#renameNestedGitRepos] ${disable ? "disabled" : "enabled"} nested git repo ${gitPath}`, + ) } catch (error) { this.log( - `failed to ${disable ? "disable" : "enable"} nested git repo ${gitPath}: ${error instanceof Error ? error.message : String(error)}`, + `[CheckpointService#renameNestedGitRepos] failed to ${disable ? "disable" : "enable"} nested git repo ${gitPath}: ${error instanceof Error ? error.message : String(error)}`, ) } } @@ -152,7 +192,7 @@ export class ShadowCheckpointService implements CheckpointService { this.shadowGitConfigWorktree = (await this.git.getConfig("core.worktree")).value || undefined } catch (error) { this.log( - `[getShadowGitConfigWorktree] failed to get core.worktree: ${error instanceof Error ? error.message : String(error)}`, + `[CheckpointService#getShadowGitConfigWorktree] failed to get core.worktree: ${error instanceof Error ? error.message : String(error)}`, ) } } @@ -162,35 +202,63 @@ export class ShadowCheckpointService implements CheckpointService { public async saveCheckpoint(message: string) { try { + this.log("[CheckpointService#saveCheckpoint] starting checkpoint save") + + if (!this.isInitialized) { + throw new Error("Shadow git repo not initialized") + } + const startTime = Date.now() await this.stageAll() const result = await this.git.commit(message) - - if (result.commit) { - const duration = Date.now() - startTime - this.log(`[saveCheckpoint] saved checkpoint ${result.commit} in ${duration}ms`) - return result - } else { - return undefined - } - } catch (error) { - this.log( - `[saveCheckpoint] failed to create checkpoint: ${error instanceof Error ? error.message : String(error)}`, - ) - + const isFirst = this._checkpoints.length === 0 + const fromHash = this._checkpoints[this._checkpoints.length - 1] ?? this.baseHash! + 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 + } catch (e) { + const error = e instanceof Error ? e : new Error(String(e)) + this.log(`[CheckpointService#saveCheckpoint] failed to create checkpoint: ${error.message}`) + this.emit("error", { type: "error", error }) throw error } } public async restoreCheckpoint(commitHash: string) { - const start = Date.now() - await this.git.clean("f", ["-d", "-f"]) - await this.git.reset(["--hard", commitHash]) - const duration = Date.now() - start - this.log(`[restoreCheckpoint] restored checkpoint ${commitHash} in ${duration}ms`) + try { + if (!this.isInitialized) { + throw new Error("Shadow git repo not initialized") + } + + const start = Date.now() + await this.git.clean("f", ["-d", "-f"]) + await this.git.reset(["--hard", commitHash]) + + // Remove all checkpoints after the specified commitHash. + const checkpointIndex = this._checkpoints.indexOf(commitHash) + + if (checkpointIndex !== -1) { + this._checkpoints = this._checkpoints.slice(0, checkpointIndex + 1) + } + + const duration = Date.now() - start + this.emit("restore", { type: "restore", commitHash, duration }) + this.log(`[CheckpointService#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.emit("error", { type: "error", error }) + throw error + } } public async getDiff({ from, to }: { from?: string; to?: string }) { + if (!this.isInitialized) { + throw new Error("Shadow git repo not initialized") + } + const result = [] if (!from) { @@ -223,6 +291,7 @@ export class ShadowCheckpointService implements CheckpointService { try { await simpleGit().version() } catch (error) { + log("[CheckpointService#create] git is not installed") throw new Error("Git must be installed to use checkpoints.") } @@ -241,9 +310,7 @@ export class ShadowCheckpointService implements CheckpointService { const gitDir = path.join(checkpointsDir, ".git") const git = simpleGit(path.dirname(gitDir)) - log(`[create] taskId = ${taskId}, workspaceDir = ${workspaceDir}, shadowDir = ${shadowDir}`) - const service = new ShadowCheckpointService(taskId, git, shadowDir, workspaceDir, log) - await service.initShadowGit() - return service + log(`[CheckpointService#create] taskId = ${taskId}, workspaceDir = ${workspaceDir}, shadowDir = ${shadowDir}`) + return new ShadowCheckpointService(taskId, git, shadowDir, workspaceDir, log) } } diff --git a/src/services/checkpoints/__tests__/LocalCheckpointService.test.ts b/src/services/checkpoints/__tests__/LocalCheckpointService.test.ts deleted file mode 100644 index 29ab7133a5..0000000000 --- a/src/services/checkpoints/__tests__/LocalCheckpointService.test.ts +++ /dev/null @@ -1,385 +0,0 @@ -// npx jest src/services/checkpoints/__tests__/LocalCheckpointService.test.ts - -import fs from "fs/promises" -import path from "path" -import os from "os" - -import { simpleGit, SimpleGit } from "simple-git" - -import { CheckpointServiceFactory } from "../CheckpointServiceFactory" -import { LocalCheckpointService } from "../LocalCheckpointService" - -const tmpDir = path.join(os.tmpdir(), "test-LocalCheckpointService") - -describe("LocalCheckpointService", () => { - const taskId = "test-task" - - let testFile: string - let service: LocalCheckpointService - - 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 { testFile } - } - - beforeEach(async () => { - const workspaceDir = path.join(tmpDir, `checkpoint-service-test-${Date.now()}`) - const repo = await initRepo({ workspaceDir }) - - testFile = repo.testFile - service = await CheckpointServiceFactory.create({ - strategy: "local", - options: { taskId, workspaceDir, log: () => {} }, - }) - }) - - afterEach(async () => { - jest.restoreAllMocks() - }) - - afterAll(async () => { - await fs.rm(tmpDir, { recursive: true, force: true }) - }) - - describe("getDiff", () => { - it("returns the correct diff between commits", async () => { - await fs.writeFile(testFile, "Ahoy, world!") - const commit1 = await service.saveCheckpoint("First checkpoint") - expect(commit1?.commit).toBeTruthy() - - await fs.writeFile(testFile, "Goodbye, world!") - const commit2 = await service.saveCheckpoint("Second checkpoint") - expect(commit2?.commit).toBeTruthy() - - const diff1 = await service.getDiff({ to: commit1!.commit }) - expect(diff1).toHaveLength(1) - expect(diff1[0].paths.relative).toBe("test.txt") - expect(diff1[0].paths.absolute).toBe(testFile) - expect(diff1[0].content.before).toBe("Hello, world!") - expect(diff1[0].content.after).toBe("Ahoy, world!") - - const diff2 = await service.getDiff({ to: commit2!.commit }) - expect(diff2).toHaveLength(1) - expect(diff2[0].paths.relative).toBe("test.txt") - expect(diff2[0].paths.absolute).toBe(testFile) - expect(diff2[0].content.before).toBe("Hello, world!") - expect(diff2[0].content.after).toBe("Goodbye, world!") - - const diff12 = await service.getDiff({ from: commit1!.commit, to: commit2!.commit }) - expect(diff12).toHaveLength(1) - expect(diff12[0].paths.relative).toBe("test.txt") - expect(diff12[0].paths.absolute).toBe(testFile) - expect(diff12[0].content.before).toBe("Ahoy, world!") - expect(diff12[0].content.after).toBe("Goodbye, world!") - }) - - it("handles new files in diff", async () => { - const newFile = path.join(service.workspaceDir, "new.txt") - await fs.writeFile(newFile, "New file content") - const commit = await service.saveCheckpoint("Add new file") - expect(commit?.commit).toBeTruthy() - - const changes = await service.getDiff({ to: commit!.commit }) - const change = changes.find((c) => c.paths.relative === "new.txt") - expect(change).toBeDefined() - expect(change?.content.before).toBe("") - expect(change?.content.after).toBe("New file content") - }) - - it("handles deleted files in diff", async () => { - const fileToDelete = path.join(service.workspaceDir, "new.txt") - await fs.writeFile(fileToDelete, "New file content") - const commit1 = await service.saveCheckpoint("Add file") - expect(commit1?.commit).toBeTruthy() - - await fs.unlink(fileToDelete) - const commit2 = await service.saveCheckpoint("Delete file") - expect(commit2?.commit).toBeTruthy() - - const changes = await service.getDiff({ from: commit1!.commit, to: commit2!.commit }) - const change = changes.find((c) => c.paths.relative === "new.txt") - expect(change).toBeDefined() - expect(change!.content.before).toBe("New file content") - expect(change!.content.after).toBe("") - }) - }) - - describe("saveCheckpoint", () => { - it("creates a checkpoint if there are pending changes", async () => { - await fs.writeFile(testFile, "Ahoy, world!") - const commit1 = await service.saveCheckpoint("First checkpoint") - expect(commit1?.commit).toBeTruthy() - const details1 = await service.git.show([commit1!.commit]) - expect(details1).toContain("-Hello, world!") - expect(details1).toContain("+Ahoy, world!") - - await fs.writeFile(testFile, "Hola, world!") - const commit2 = await service.saveCheckpoint("Second checkpoint") - expect(commit2?.commit).toBeTruthy() - const details2 = await service.git.show([commit2!.commit]) - expect(details2).toContain("-Hello, world!") - expect(details2).toContain("+Hola, world!") - - // Switch to checkpoint 1. - await service.restoreCheckpoint(commit1!.commit) - expect(await fs.readFile(testFile, "utf-8")).toBe("Ahoy, world!") - - // Switch to checkpoint 2. - await service.restoreCheckpoint(commit2!.commit) - expect(await fs.readFile(testFile, "utf-8")).toBe("Hola, world!") - - // Switch back to initial commit. - expect(service.baseHash).toBeTruthy() - await service.restoreCheckpoint(service.baseHash!) - expect(await fs.readFile(testFile, "utf-8")).toBe("Hello, world!") - }) - - it("preserves workspace and index state after saving checkpoint", async () => { - // Create three files with different states: staged, unstaged, and mixed. - const unstagedFile = path.join(service.workspaceDir, "unstaged.txt") - const stagedFile = path.join(service.workspaceDir, "staged.txt") - const mixedFile = path.join(service.workspaceDir, "mixed.txt") - - await fs.writeFile(unstagedFile, "Initial unstaged") - await fs.writeFile(stagedFile, "Initial staged") - await fs.writeFile(mixedFile, "Initial mixed") - await service.git.add(["."]) - const result = await service.git.commit("Add initial files") - expect(result?.commit).toBeTruthy() - - await fs.writeFile(unstagedFile, "Modified unstaged") - - await fs.writeFile(stagedFile, "Modified staged") - await service.git.add([stagedFile]) - - await fs.writeFile(mixedFile, "Modified mixed - staged") - await service.git.add([mixedFile]) - await fs.writeFile(mixedFile, "Modified mixed - unstaged") - - // Save checkpoint. - const commit = await service.saveCheckpoint("Test checkpoint") - expect(commit?.commit).toBeTruthy() - - // Verify workspace state is preserved. - const status = await service.git.status() - - // All files should be modified. - expect(status.modified).toContain("unstaged.txt") - expect(status.modified).toContain("staged.txt") - expect(status.modified).toContain("mixed.txt") - - // Only staged and mixed files should be staged. - expect(status.staged).not.toContain("unstaged.txt") - expect(status.staged).toContain("staged.txt") - expect(status.staged).toContain("mixed.txt") - - // Verify file contents. - expect(await fs.readFile(unstagedFile, "utf-8")).toBe("Modified unstaged") - expect(await fs.readFile(stagedFile, "utf-8")).toBe("Modified staged") - expect(await fs.readFile(mixedFile, "utf-8")).toBe("Modified mixed - unstaged") - - // Verify staged changes (--cached shows only staged changes). - const stagedDiff = await service.git.diff(["--cached", "mixed.txt"]) - expect(stagedDiff).toContain("-Initial mixed") - expect(stagedDiff).toContain("+Modified mixed - staged") - - // Verify unstaged changes (shows working directory changes). - const unstagedDiff = await service.git.diff(["mixed.txt"]) - expect(unstagedDiff).toContain("-Modified mixed - staged") - expect(unstagedDiff).toContain("+Modified mixed - unstaged") - }) - - it("does not create a checkpoint if there are no pending changes", async () => { - const commit0 = await service.saveCheckpoint("Zeroth checkpoint") - expect(commit0?.commit).toBeFalsy() - - await fs.writeFile(testFile, "Ahoy, world!") - const commit1 = await service.saveCheckpoint("First checkpoint") - expect(commit1?.commit).toBeTruthy() - - const commit2 = await service.saveCheckpoint("Second checkpoint") - expect(commit2?.commit).toBeFalsy() - }) - - it("includes untracked files in checkpoints", async () => { - // Create an untracked file. - const untrackedFile = path.join(service.workspaceDir, "untracked.txt") - await fs.writeFile(untrackedFile, "I am untracked!") - - // Save a checkpoint with the untracked file. - const commit1 = await service.saveCheckpoint("Checkpoint with untracked file") - expect(commit1?.commit).toBeTruthy() - - // Verify the untracked file was included in the checkpoint. - const details = await service.git.show([commit1!.commit]) - expect(details).toContain("+I am untracked!") - - // Create another checkpoint with a different state. - await fs.writeFile(testFile, "Changed tracked file") - const commit2 = await service.saveCheckpoint("Second checkpoint") - expect(commit2?.commit).toBeTruthy() - - // Restore first checkpoint and verify untracked file is preserved. - await service.restoreCheckpoint(commit1!.commit) - expect(await fs.readFile(untrackedFile, "utf-8")).toBe("I am untracked!") - expect(await fs.readFile(testFile, "utf-8")).toBe("Hello, world!") - - // Restore second checkpoint and verify untracked file remains (since - // restore preserves untracked files) - await service.restoreCheckpoint(commit2!.commit) - expect(await fs.readFile(untrackedFile, "utf-8")).toBe("I am untracked!") - expect(await fs.readFile(testFile, "utf-8")).toBe("Changed tracked file") - }) - - it("throws if we're on the wrong branch", async () => { - // Create and switch to a feature branch. - const currentBranch = await service.git.revparse(["--abbrev-ref", "HEAD"]) - await service.git.checkoutBranch("feature", currentBranch) - - // Attempt to save checkpoint from feature branch. - await expect(service.saveCheckpoint("test")).rejects.toThrow( - `Git branch mismatch: expected '${currentBranch}' but found 'feature'`, - ) - - // Attempt to restore checkpoint from feature branch. - expect(service.baseHash).toBeTruthy() - - await expect(service.restoreCheckpoint(service.baseHash!)).rejects.toThrow( - `Git branch mismatch: expected '${currentBranch}' but found 'feature'`, - ) - }) - - it("cleans up staged files if a commit fails", async () => { - await fs.writeFile(testFile, "Changed content") - - // Mock git commit to simulate failure. - jest.spyOn(service.git, "commit").mockRejectedValue(new Error("Simulated commit failure")) - - // Attempt to save checkpoint. - await expect(service.saveCheckpoint("test")).rejects.toThrow("Simulated commit failure") - - // Verify files are unstaged. - const status = await service.git.status() - expect(status.staged).toHaveLength(0) - }) - - it("handles file deletions correctly", async () => { - await fs.writeFile(testFile, "I am tracked!") - const untrackedFile = path.join(service.workspaceDir, "new.txt") - await fs.writeFile(untrackedFile, "I am untracked!") - const commit1 = await service.saveCheckpoint("First checkpoint") - expect(commit1?.commit).toBeTruthy() - - await fs.unlink(testFile) - await fs.unlink(untrackedFile) - const commit2 = await service.saveCheckpoint("Second checkpoint") - expect(commit2?.commit).toBeTruthy() - - // Verify files are gone. - await expect(fs.readFile(testFile, "utf-8")).rejects.toThrow() - await expect(fs.readFile(untrackedFile, "utf-8")).rejects.toThrow() - - // Restore first checkpoint. - await service.restoreCheckpoint(commit1!.commit) - expect(await fs.readFile(testFile, "utf-8")).toBe("I am tracked!") - expect(await fs.readFile(untrackedFile, "utf-8")).toBe("I am untracked!") - - // Restore second checkpoint. - await service.restoreCheckpoint(commit2!.commit) - await expect(fs.readFile(testFile, "utf-8")).rejects.toThrow() - await expect(fs.readFile(untrackedFile, "utf-8")).rejects.toThrow() - }) - }) - - describe("create", () => { - it("initializes a git repository if one does not already exist", async () => { - const workspaceDir = path.join(tmpDir, `checkpoint-service-test2-${Date.now()}`) - await fs.mkdir(workspaceDir) - const newTestFile = path.join(workspaceDir, "test.txt") - await fs.writeFile(newTestFile, "Hello, world!") - - // Ensure the git repository was initialized. - const gitDir = path.join(workspaceDir, ".git") - await expect(fs.stat(gitDir)).rejects.toThrow() - const newService = await LocalCheckpointService.create({ taskId, workspaceDir, log: () => {} }) - expect(await fs.stat(gitDir)).toBeTruthy() - - // Save a checkpoint: Hello, world! - const commit1 = await newService.saveCheckpoint("Hello, world!") - expect(commit1?.commit).toBeTruthy() - expect(await fs.readFile(newTestFile, "utf-8")).toBe("Hello, world!") - - // Restore initial commit; the file should no longer exist. - expect(newService.baseHash).toBeTruthy() - await newService.restoreCheckpoint(newService.baseHash!) - await expect(fs.access(newTestFile)).rejects.toThrow() - - // Restore to checkpoint 1; the file should now exist. - await newService.restoreCheckpoint(commit1!.commit) - expect(await fs.readFile(newTestFile, "utf-8")).toBe("Hello, world!") - - // Save a new checkpoint: Ahoy, world! - await fs.writeFile(newTestFile, "Ahoy, world!") - const commit2 = await newService.saveCheckpoint("Ahoy, world!") - expect(commit2?.commit).toBeTruthy() - expect(await fs.readFile(newTestFile, "utf-8")).toBe("Ahoy, world!") - - // Restore "Hello, world!" - await newService.restoreCheckpoint(commit1!.commit) - expect(await fs.readFile(newTestFile, "utf-8")).toBe("Hello, world!") - - // Restore "Ahoy, world!" - await newService.restoreCheckpoint(commit2!.commit) - expect(await fs.readFile(newTestFile, "utf-8")).toBe("Ahoy, world!") - - // Restore initial commit. - expect(newService.baseHash).toBeTruthy() - await newService.restoreCheckpoint(newService.baseHash!) - await expect(fs.access(newTestFile)).rejects.toThrow() - - await fs.rm(newService.workspaceDir, { recursive: true, force: true }) - }) - - it("respects existing git user configuration", async () => { - const workspaceDir = path.join(tmpDir, `checkpoint-service-test-config2-${Date.now()}`) - const userName = "Custom User" - const userEmail = "custom@example.com" - await initRepo({ workspaceDir, userName, userEmail }) - - const newService = await LocalCheckpointService.create({ taskId, workspaceDir, log: () => {} }) - - expect((await newService.git.getConfig("user.name")).value).toBe(userName) - expect((await newService.git.getConfig("user.email")).value).toBe(userEmail) - - await fs.rm(workspaceDir, { recursive: true, force: true }) - }) - }) -}) diff --git a/src/services/checkpoints/__tests__/ShadowCheckpointService.test.ts b/src/services/checkpoints/__tests__/ShadowCheckpointService.test.ts index 5bcc5d444c..138e1130e8 100644 --- a/src/services/checkpoints/__tests__/ShadowCheckpointService.test.ts +++ b/src/services/checkpoints/__tests__/ShadowCheckpointService.test.ts @@ -3,11 +3,11 @@ import fs from "fs/promises" import path from "path" import os from "os" +import { EventEmitter } from "events" import { simpleGit, SimpleGit } from "simple-git" import { ShadowCheckpointService } from "../ShadowCheckpointService" -import { CheckpointServiceFactory } from "../CheckpointServiceFactory" jest.mock("globby", () => ({ globby: jest.fn().mockResolvedValue([]), @@ -63,13 +63,10 @@ describe("ShadowCheckpointService", () => { const repo = await initRepo({ workspaceDir }) testFile = repo.testFile - - service = await CheckpointServiceFactory.create({ - strategy: "shadow", - options: { taskId, shadowDir, workspaceDir, log: () => {} }, - }) - workspaceGit = repo.git + + service = await ShadowCheckpointService.create({ taskId, shadowDir, workspaceDir, log: () => {} }) + await service.initShadowGit() }) afterEach(async () => { @@ -313,6 +310,7 @@ describe("ShadowCheckpointService", () => { 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() expect(await fs.stat(gitDir)).toBeTruthy() // Save a new checkpoint: Ahoy, world! @@ -333,4 +331,168 @@ describe("ShadowCheckpointService", () => { await fs.rm(newService.workspaceDir, { recursive: true, force: true }) }) }) + + describe("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()}`) + await fs.mkdir(workspaceDir, { recursive: true }) + + const newTestFile = path.join(workspaceDir, "test.txt") + await fs.writeFile(newTestFile, "Testing events!") + + // Create a mock implementation of emit to track events. + 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: () => {} }) + await newService.initShadowGit() + + // Find the initialize event in the emit calls. + let initializeEvent = null + + for (let i = 0; i < emitSpy.mock.calls.length; i++) { + const call = emitSpy.mock.calls[i] + + if (call[0] === "initialize") { + initializeEvent = call[1] + break + } + } + + // Restore the spy. + emitSpy.mockRestore() + + // Verify the event was emitted with the correct data. + expect(initializeEvent).not.toBeNull() + expect(initializeEvent.type).toBe("initialize") + expect(initializeEvent.workspaceDir).toBe(workspaceDir) + expect(initializeEvent.baseHash).toBeTruthy() + expect(typeof initializeEvent.created).toBe("boolean") + expect(typeof initializeEvent.duration).toBe("number") + + // Verify the event was emitted with the correct data. + expect(initializeEvent).not.toBeNull() + expect(initializeEvent.type).toBe("initialize") + expect(initializeEvent.workspaceDir).toBe(workspaceDir) + expect(initializeEvent.baseHash).toBeTruthy() + expect(typeof initializeEvent.created).toBe("boolean") + expect(typeof initializeEvent.duration).toBe("number") + + // Clean up. + await fs.rm(shadowDir, { recursive: true, force: true }) + await fs.rm(workspaceDir, { recursive: true, force: true }) + }) + + it("emits checkpoint event when saving checkpoint", async () => { + const checkpointHandler = jest.fn() + service.on("checkpoint", checkpointHandler) + + await fs.writeFile(testFile, "Changed content for checkpoint event test") + const result = await service.saveCheckpoint("Test checkpoint event") + expect(result?.commit).toBeDefined() + + expect(checkpointHandler).toHaveBeenCalledTimes(1) + const eventData = checkpointHandler.mock.calls[0][0] + expect(eventData.type).toBe("checkpoint") + expect(eventData.toHash).toBeDefined() + expect(eventData.toHash).toBe(result!.commit) + expect(typeof eventData.duration).toBe("number") + }) + + it("emits restore event when restoring checkpoint", async () => { + // First create a checkpoint to restore. + await fs.writeFile(testFile, "Content for restore test") + const commit = await service.saveCheckpoint("Checkpoint for restore test") + expect(commit?.commit).toBeTruthy() + + // Change the file again. + await fs.writeFile(testFile, "Changed after checkpoint") + + // Setup restore event listener. + const restoreHandler = jest.fn() + service.on("restore", restoreHandler) + + // Restore the checkpoint. + await service.restoreCheckpoint(commit!.commit) + + // Verify the event was emitted. + expect(restoreHandler).toHaveBeenCalledTimes(1) + const eventData = restoreHandler.mock.calls[0][0] + expect(eventData.type).toBe("restore") + expect(eventData.commitHash).toBe(commit!.commit) + expect(typeof eventData.duration).toBe("number") + + // Verify the file was actually restored. + expect(await fs.readFile(testFile, "utf-8")).toBe("Content for restore test") + }) + + it("emits error event when an error occurs", async () => { + const errorHandler = jest.fn() + service.on("error", errorHandler) + + // Force an error by providing an invalid commit hash. + const invalidCommitHash = "invalid-commit-hash" + + // Try to restore an invalid checkpoint. + try { + await service.restoreCheckpoint(invalidCommitHash) + } catch (error) { + // Expected to throw, we're testing the event emission. + } + + // Verify the error event was emitted. + expect(errorHandler).toHaveBeenCalledTimes(1) + const eventData = errorHandler.mock.calls[0][0] + expect(eventData.type).toBe("error") + expect(eventData.error).toBeInstanceOf(Error) + }) + + it("supports multiple event listeners for the same event", async () => { + const checkpointHandler1 = jest.fn() + const checkpointHandler2 = jest.fn() + + service.on("checkpoint", checkpointHandler1) + service.on("checkpoint", checkpointHandler2) + + await fs.writeFile(testFile, "Content for multiple listeners test") + const result = await service.saveCheckpoint("Testing multiple listeners") + + // Verify both handlers were called with the same event data. + expect(checkpointHandler1).toHaveBeenCalledTimes(1) + expect(checkpointHandler2).toHaveBeenCalledTimes(1) + + const eventData1 = checkpointHandler1.mock.calls[0][0] + const eventData2 = checkpointHandler2.mock.calls[0][0] + + expect(eventData1).toEqual(eventData2) + expect(eventData1.type).toBe("checkpoint") + expect(eventData1.toHash).toBe(result?.commit) + }) + + it("allows removing event listeners", async () => { + const checkpointHandler = jest.fn() + + // Add the listener. + service.on("checkpoint", checkpointHandler) + + // Make a change and save a checkpoint. + await fs.writeFile(testFile, "Content for remove listener test - part 1") + await service.saveCheckpoint("Testing listener - part 1") + + // Verify handler was called. + expect(checkpointHandler).toHaveBeenCalledTimes(1) + checkpointHandler.mockClear() + + // Remove the listener. + service.off("checkpoint", checkpointHandler) + + // Make another change and save a checkpoint. + await fs.writeFile(testFile, "Content for remove listener test - part 2") + await service.saveCheckpoint("Testing listener - part 2") + + // Verify handler was not called after being removed. + expect(checkpointHandler).not.toHaveBeenCalled() + }) + }) }) diff --git a/src/services/checkpoints/index.ts b/src/services/checkpoints/index.ts deleted file mode 100644 index e62b6119ce..0000000000 --- a/src/services/checkpoints/index.ts +++ /dev/null @@ -1,2 +0,0 @@ -export * from "./types" -export * from "./CheckpointServiceFactory" diff --git a/src/services/checkpoints/types.ts b/src/services/checkpoints/types.ts index 50843abbd3..3ae9457fe8 100644 --- a/src/services/checkpoints/types.ts +++ b/src/services/checkpoints/types.ts @@ -1,3 +1,4 @@ +import EventEmitter from "events" import { CommitResult } from "simple-git" export type CheckpointResult = Partial & Pick @@ -13,15 +14,12 @@ export type CheckpointDiff = { } } -export type CheckpointStrategy = "local" | "shadow" - export interface CheckpointService { saveCheckpoint(message: string): Promise restoreCheckpoint(commit: string): Promise getDiff(range: { from?: string; to?: string }): Promise workspaceDir: string baseHash?: string - strategy: CheckpointStrategy version: number } @@ -30,3 +28,38 @@ export interface CheckpointServiceOptions { workspaceDir: string log?: (message: string) => void } + +/** + * EventEmitter + */ + +export interface CheckpointEventMap { + initialize: { type: "initialize"; workspaceDir: string; baseHash: string; created: boolean; duration: number } + checkpoint: { + type: "checkpoint" + isFirst: boolean + fromHash: string + toHash: string + duration: number + } + 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/webview-ui/src/components/chat/checkpoints/CheckpointSaved.tsx b/webview-ui/src/components/chat/checkpoints/CheckpointSaved.tsx index c15bbd102c..5d3475c4d0 100644 --- a/webview-ui/src/components/chat/checkpoints/CheckpointSaved.tsx +++ b/webview-ui/src/components/chat/checkpoints/CheckpointSaved.tsx @@ -3,8 +3,6 @@ import { useMemo } from "react" import { CheckpointMenu } from "./CheckpointMenu" import { checkpointSchema } from "./schema" -const REQUIRED_VERSION = 1 - type CheckpointSavedProps = { ts: number commitHash: string @@ -22,7 +20,7 @@ export const CheckpointSaved = ({ checkpoint, ...props }: CheckpointSavedProps) const result = checkpointSchema.safeParse(checkpoint) - if (!result.success || result.data.version < REQUIRED_VERSION) { + if (!result.success) { return undefined } diff --git a/webview-ui/src/components/chat/checkpoints/schema.ts b/webview-ui/src/components/chat/checkpoints/schema.ts index 7f097966b8..4acd32a6ab 100644 --- a/webview-ui/src/components/chat/checkpoints/schema.ts +++ b/webview-ui/src/components/chat/checkpoints/schema.ts @@ -4,8 +4,6 @@ export const checkpointSchema = z.object({ isFirst: z.boolean(), from: z.string(), to: z.string(), - strategy: z.enum(["local", "shadow"]), - version: z.number(), }) export type Checkpoint = z.infer