diff --git a/src/services/checkpoints/CheckpointService.ts b/src/services/checkpoints/CheckpointService.ts new file mode 100644 index 0000000000..af1d438e54 --- /dev/null +++ b/src/services/checkpoints/CheckpointService.ts @@ -0,0 +1,317 @@ +import fs from "fs/promises" +import { existsSync } from "fs" +import path from "path" + +import debug from "debug" +import simpleGit, { SimpleGit, CleanOptions } from "simple-git" + +if (process.env.NODE_ENV !== "test") { + debug.enable("simple-git") +} + +export interface Checkpoint { + hash: string + message: string + timestamp?: Date +} + +export type CheckpointServiceOptions = { + taskId: string + git?: SimpleGit + baseDir: string + log?: (message: string) => void +} + +/** + * 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: + * - Current changes are stashed (including untracked files). + * - The hidden branch is reset to match main. + * - Stashed changes are applied and committed as a checkpoint on the hidden + * branch. + * - We return to the main branch with the original state restored. + * + * 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. + * + * 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 CheckpointService { + constructor( + public readonly taskId: string, + private readonly git: SimpleGit, + public readonly baseDir: string, + public readonly mainBranch: string, + public readonly baseCommitHash: string, + public readonly hiddenBranch: string, + private readonly log: (message: string) => void, + ) {} + + private async pushStash() { + const status = await this.git.status() + + if (status.files.length > 0) { + await this.git.stash(["-u"]) // Includes tracked and untracked files. + return true + } + + return false + } + + private async applyStash() { + const stashList = await this.git.stashList() + + if (stashList.all.length > 0) { + await this.git.stash(["apply"]) // Applies the most recent stash only. + return true + } + + return false + } + + private async popStash() { + const stashList = await this.git.stashList() + + if (stashList.all.length > 0) { + await this.git.stash(["pop", "--index"]) // Pops the most recent stash only. + return true + } + + return false + } + + 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.baseCommitHash + } + + 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.baseDir, relPath) + + // If modified both before and after will generate content. + // If added only after will generate content. + // If deleted only before will generate content. + let beforeContent = "" + let afterContent = "" + + try { + beforeContent = await this.git.show([`${from}:${relPath}`]) + } catch (err) { + // File doesn't exist in older commit. + } + + try { + afterContent = await this.git.show([`${to}:${relPath}`]) + } catch (err) { + // File doesn't exist in newer commit. + } + + result.push({ + paths: { relative: relPath, absolute: absPath }, + content: { before: beforeContent, after: afterContent }, + }) + } + + return result + } + + public async saveCheckpoint(message: string) { + await this.ensureBranch(this.mainBranch) + + // Attempt to stash pending changes (including untracked files). + const pendingChanges = await this.pushStash() + + // Get the latest commit on the hidden branch before we reset it. + const latestHash = await this.git.revparse([this.hiddenBranch]) + + // Check if there is any diff relative to the latest commit. + if (!pendingChanges) { + const diff = await this.git.diff([latestHash]) + + if (!diff) { + this.log(`[saveCheckpoint] No changes detected, giving up`) + return undefined + } + } + + await this.git.checkout(this.hiddenBranch) + + const reset = async () => { + await this.git.reset(["HEAD", "."]) + await this.git.clean([CleanOptions.FORCE, CleanOptions.RECURSIVE]) + await this.git.reset(["--hard", latestHash]) + await this.git.checkout(this.mainBranch) + await this.popStash() + } + + try { + // Reset hidden branch to match main and apply the pending changes. + await this.git.reset(["--hard", this.mainBranch]) + + if (pendingChanges) { + await this.applyStash() + } + + // Using "-A" ensures that deletions are staged as well. + await this.git.add(["-A"]) + const diff = await this.git.diff([latestHash]) + + if (!diff) { + this.log(`[saveCheckpoint] No changes detected, resetting and giving up`) + await reset() + return undefined + } + + // Otherwise, commit the changes. + const status = await this.git.status() + this.log(`[saveCheckpoint] Changes detected, committing ${JSON.stringify(status)}`) + + // Allow empty commits in order to correctly handle deletion of + // untracked files (see unit tests for an example of this). + // Additionally, skip pre-commit hooks so that they don't slow + // things down or tamper with the contents of the commit. + const commit = await this.git.commit(message, undefined, { + "--allow-empty": null, + "--no-verify": null, + }) + + await this.git.checkout(this.mainBranch) + + if (pendingChanges) { + await this.popStash() + } + + return commit + } catch (err) { + this.log(`[saveCheckpoint] Failed to save checkpoint: ${err instanceof Error ? err.message : String(err)}`) + + // If we're not on the main branch then we need to trigger a reset + // to return to the main branch and restore it's previous state. + const currentBranch = await this.git.revparse(["--abbrev-ref", "HEAD"]) + + if (currentBranch.trim() !== this.mainBranch) { + await reset() + } + + throw err + } + } + + public async restoreCheckpoint(commitHash: string) { + await this.ensureBranch(this.mainBranch) + await this.git.clean([CleanOptions.FORCE, CleanOptions.RECURSIVE]) + await this.git.raw(["restore", "--source", commitHash, "--worktree", "--", "."]) + } + + public static async create({ taskId, git, baseDir, log = console.log }: CheckpointServiceOptions) { + git = + git || + simpleGit({ + baseDir, + binary: "git", + maxConcurrentProcesses: 1, + config: [], + trimmed: true, + }) + + 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 (!baseDir || !existsSync(baseDir)) { + throw new Error(`Base directory is not set or does not exist.`) + } + + const { currentBranch, currentSha, hiddenBranch } = await CheckpointService.initRepo({ + taskId, + git, + baseDir, + log, + }) + + log( + `[CheckpointService] taskId = ${taskId}, baseDir = ${baseDir}, currentBranch = ${currentBranch}, currentSha = ${currentSha}, hiddenBranch = ${hiddenBranch}`, + ) + return new CheckpointService(taskId, git, baseDir, currentBranch, currentSha, hiddenBranch, log) + } + + private static async initRepo({ taskId, git, baseDir, log }: Required) { + const isExistingRepo = existsSync(path.join(baseDir, ".git")) + + if (!isExistingRepo) { + await git.init() + log(`[initRepo] Initialized new Git repository at ${baseDir}`) + } + + await git.addConfig("user.name", "Roo Code") + await git.addConfig("user.email", "support@roocode.com") + + 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(baseDir, ".gitkeep"), "") + await git.add(".") + 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 = `roo-code-checkpoints-${taskId}` + const branchSummary = await git.branch() + + if (!branchSummary.all.includes(hiddenBranch)) { + await git.checkoutBranch(hiddenBranch, currentBranch) // git checkout -b + await git.checkout(currentBranch) // git checkout + } + + return { currentBranch, currentSha, hiddenBranch } + } +} diff --git a/src/services/checkpoints/__tests__/CheckpointService.test.ts b/src/services/checkpoints/__tests__/CheckpointService.test.ts new file mode 100644 index 0000000000..cd33a5dc7c --- /dev/null +++ b/src/services/checkpoints/__tests__/CheckpointService.test.ts @@ -0,0 +1,337 @@ +// npx jest src/services/checkpoints/__tests__/CheckpointService.test.ts + +import fs from "fs/promises" +import path from "path" +import os from "os" + +import { simpleGit, SimpleGit } from "simple-git" + +import { CheckpointService } from "../CheckpointService" + +describe("CheckpointService", () => { + const taskId = "test-task" + let git: SimpleGit + let testFile: string + let service: CheckpointService + + beforeEach(async () => { + // Create a temporary directory for testing. + const baseDir = path.join(os.tmpdir(), `checkpoint-service-test-${Date.now()}`) + await fs.mkdir(baseDir) + + // Initialize git repo. + git = simpleGit(baseDir) + await git.init() + await git.addConfig("user.name", "Roo Code") + await git.addConfig("user.email", "support@roo.vet") + + // Create test file. + testFile = path.join(baseDir, "test.txt") + await fs.writeFile(testFile, "Hello, world!") + + // Create initial commit. + await git.add(".") + await git.commit("Initial commit")! + + // Create service instance. + const log = () => {} + service = await CheckpointService.create({ taskId, git, baseDir, log }) + }) + + afterEach(async () => { + await fs.rm(service.baseDir, { recursive: true, force: true }) + jest.restoreAllMocks() + }) + + 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.baseDir, "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.baseDir, "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 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 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. + await service.restoreCheckpoint(service.baseCommitHash) + 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.baseDir, "unstaged.txt") + const stagedFile = path.join(service.baseDir, "staged.txt") + const mixedFile = path.join(service.baseDir, "mixed.txt") + + await fs.writeFile(unstagedFile, "Initial unstaged") + await fs.writeFile(stagedFile, "Initial staged") + await fs.writeFile(mixedFile, "Initial mixed") + await git.add(["."]) + const result = await git.commit("Add initial files") + expect(result?.commit).toBeTruthy() + + await fs.writeFile(unstagedFile, "Modified unstaged") + + await fs.writeFile(stagedFile, "Modified staged") + await git.add([stagedFile]) + + await fs.writeFile(mixedFile, "Modified mixed - staged") + await 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 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 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 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 () => { + await fs.writeFile(testFile, "Ahoy, world!") + const commit = await service.saveCheckpoint("First checkpoint") + expect(commit?.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.baseDir, "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 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. + await git.checkoutBranch("feature", service.mainBranch) + + // Attempt to save checkpoint from feature branch. + await expect(service.saveCheckpoint("test")).rejects.toThrow( + `Git branch mismatch: expected '${service.mainBranch}' but found 'feature'`, + ) + + // Attempt to restore checkpoint from feature branch. + await expect(service.restoreCheckpoint(service.baseCommitHash)).rejects.toThrow( + `Git branch mismatch: expected '${service.mainBranch}' 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(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 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.baseDir, "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 baseDir = path.join(os.tmpdir(), `checkpoint-service-test2-${Date.now()}`) + await fs.mkdir(baseDir) + const newTestFile = path.join(baseDir, "test.txt") + + const newGit = simpleGit(baseDir) + const initSpy = jest.spyOn(newGit, "init") + const newService = await CheckpointService.create({ taskId, git: newGit, baseDir, log: () => {} }) + + // Ensure the git repository was initialized. + expect(initSpy).toHaveBeenCalled() + + // Save a checkpoint: Hello, world! + await fs.writeFile(newTestFile, "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. + await newService.restoreCheckpoint(newService.baseCommitHash) + 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. + await newService.restoreCheckpoint(newService.baseCommitHash) + await expect(fs.access(newTestFile)).rejects.toThrow() + + await fs.rm(newService.baseDir, { recursive: true, force: true }) + }) + }) +})