Worktree checkpoints

src/services/checkpoints/ShadowCheckpointService.ts

@@ -20,12 +20,12 @@ export abstract class ShadowCheckpointService extends EventEmitter {
 	public readonly checkpointsDir: string
 	public readonly workspaceDir: string
 
+	protected configDir: string
+	protected readonly log: (message: string) => void
+
 	protected _checkpoints: string[] = []
 	protected _baseHash?: string
-
-	protected readonly dotGitDir: string
 	protected git?: SimpleGit
-	protected readonly log: (message: string) => void
 	protected shadowGitConfigWorktree?: string
 
 	public get baseHash() {
@@ -56,8 +56,8 @@ export abstract class ShadowCheckpointService extends EventEmitter {
 		this.taskId = taskId
 		this.checkpointsDir = checkpointsDir
 		this.workspaceDir = workspaceDir
+		this.configDir = path.join(this.checkpointsDir, ".git")
 
-		this.dotGitDir = path.join(this.checkpointsDir, ".git")
 		this.log = log
 	}
 
@@ -66,47 +66,16 @@ export abstract class ShadowCheckpointService extends EventEmitter {
 			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}`)
-
-		let created = false
 		const startTime = Date.now()
-
-		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(
-					`Checkpoints can only be used in the original workspace: ${worktree} !== ${this.workspaceDir}`,
-				)
-			}
-
-			await this.writeExcludeFile()
-			this.baseHash = await git.revparse(["HEAD"])
-		} else {
-			this.log(`[${this.constructor.name}#initShadowGit] creating shadow git repo at ${this.checkpointsDir}`)
-			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", "[email protected]")
-			await this.writeExcludeFile()
-			await this.stageAll(git)
-			const { commit } = await git.commit("initial commit", { "--allow-empty": null })
-			this.baseHash = commit
-			created = true
-		}
-
+		const { git, baseHash, created } = await this.initializeShadowRepo()
 		const duration = Date.now() - startTime
 
 		this.log(
-			`[${this.constructor.name}#initShadowGit] initialized shadow repo with base commit ${this.baseHash} in ${duration}ms`,
+			`[${this.constructor.name}#initShadowGit] initialized shadow repo with base commit ${baseHash} in ${duration}ms`,
 		)
 
 		this.git = git
+		this.baseHash = baseHash
 
 		await onInit?.()
 
@@ -121,22 +90,68 @@ export abstract class ShadowCheckpointService extends EventEmitter {
 		return { created, duration }
 	}
 
+	protected isShadowRepoAvailable() {
+		return fileExistsAtPath(this.configDir)
+	}
+
+	protected async initializeShadowRepo() {
+		await fs.mkdir(this.checkpointsDir, { recursive: true })
+		const git = simpleGit(this.checkpointsDir)
+		const gitVersion = await git.version()
+		this.log(`[${this.constructor.name}#initializeShadowRepo] git = ${gitVersion}`)
+		const exists = await this.isShadowRepoAvailable()
+		console.log(`[${this.constructor.name}#initializeShadowRepo] exists = ${exists} [${this.configDir}]`)
+		const baseHash = exists ? await this.checkShadowRepo(git) : await this.createShadowRepo(git)
+		return { git, baseHash, created: !exists }
+	}
+
+	protected async checkShadowRepo(git: SimpleGit) {
+		this.log(`[${this.constructor.name}#checkShadowRepo] checking existing shadow repo at ${this.configDir}`)
+		const worktree = await this.getShadowGitConfigWorktree(git)
+
+		if (worktree !== this.workspaceDir) {
+			throw new Error(
+				`Checkpoints can only be used in the original workspace: ${worktree} !== ${this.workspaceDir}`,
+			)
+		}
+
+		await this.writeExcludeFile()
+		return await git.revparse(["HEAD"])
+	}
+
+	protected async createShadowRepo(git: SimpleGit) {
+		this.log(`[${this.constructor.name}#createShadowRepo] creating new shadow repo at ${this.checkpointsDir}`)
+		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", "[email protected]")
+		await this.writeExcludeFile()
+		await this.stageAll(git)
+		const result = await git.commit("initial commit", { "--allow-empty": null })
+		return result.commit
+	}
+
 	// Add basic excludes directly in git config, while respecting any
 	// .gitignore in the workspace.
 	// .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.
 	protected async writeExcludeFile() {
-		await fs.mkdir(path.join(this.dotGitDir, "info"), { recursive: true })
+		await fs.mkdir(path.join(this.configDir, "info"), { recursive: true })
 		const patterns = await getExcludePatterns(this.workspaceDir)
-		await fs.writeFile(path.join(this.dotGitDir, "info", "exclude"), patterns.join("\n"))
+		await fs.writeFile(path.join(this.configDir, "info", "exclude"), patterns.join("\n"))
+	}
+
+	protected stagePath(git: SimpleGit, path: string) {
+		return git.add(path)
 	}
 
-	private async stageAll(git: SimpleGit) {
+	protected async stageAll(git: SimpleGit) {
 		await this.renameNestedGitRepos(true)
 
 		try {
-			await git.add(".")
+			await this.stagePath(git, ".")
 		} catch (error) {
 			this.log(
 				`[${this.constructor.name}#stageAll] failed to add files to git: ${error instanceof Error ? error.message : String(error)}`,
@@ -149,7 +164,7 @@ export abstract class ShadowCheckpointService extends EventEmitter {
 	// Since we use git to track checkpoints, we need to temporarily disable
 	// nested git repos to work around git's requirement of using submodules for
 	// nested repos.
-	private async renameNestedGitRepos(disable: boolean) {
+	protected async renameNestedGitRepos(disable: boolean) {
 		// Find all .git directories that are not at the root level.
 		const gitPaths = await globby("**/.git" + (disable ? "" : GIT_DISABLED_SUFFIX), {
 			cwd: this.workspaceDir,
@@ -186,7 +201,7 @@ export abstract class ShadowCheckpointService extends EventEmitter {
 		}
 	}
 
-	private async getShadowGitConfigWorktree(git: SimpleGit) {
+	protected async getShadowGitConfigWorktree(git: SimpleGit) {
 		if (!this.shadowGitConfigWorktree) {
 			try {
 				this.shadowGitConfigWorktree = (await git.getConfig("core.worktree")).value || undefined

src/services/checkpoints/WorktreeCheckpointService.ts

@@ -0,0 +1,75 @@
+import fs from "fs/promises"
+import * as path from "path"
+
+import simpleGit, { SimpleGit } from "simple-git"
+
+import { fileExistsAtPath } from "../../utils/fs"
+
+import { CheckpointServiceOptions } from "./types"
+import { ShadowCheckpointService } from "./ShadowCheckpointService"
+
+export class WorktreeCheckpointService extends ShadowCheckpointService {
+	private readonly worktreeDir: string
+
+	constructor(taskId: string, checkpointsDir: string, workspaceDir: string, log = console.log) {
+		super(taskId, checkpointsDir, workspaceDir, log)
+		this.configDir = this.checkpointsDir
+		this.worktreeDir = path.join(path.dirname(this.checkpointsDir), taskId)
+	}
+
+	protected override isShadowRepoAvailable() {
+		return fileExistsAtPath(path.join(this.configDir, "workspace"))
+	}
+
+	protected override async initializeShadowRepo() {
+		const result = await super.initializeShadowRepo()
+		const git = simpleGit(this.worktreeDir)
+		return { ...result, git }
+	}
+
+	protected override async createShadowRepo(git: SimpleGit) {
+		this.log(`[${this.constructor.name}#createShadowRepo] creating bare shadow git repo at ${this.configDir}`)
+		await git.init(["--bare"])
+
+		await fs.writeFile(path.join(this.configDir, "workspace"), this.workspaceDir)
+
+		await git.addConfig("commit.gpgSign", "false") // Disable commit signing for shadow repo.
+		await git.addConfig("user.name", "Roo Code")
+		await git.addConfig("user.email", "[email protected]")
+
+		await git.raw(["worktree", "add", "--orphan", this.worktreeDir])
+
+		const worktreeGit = simpleGit(this.worktreeDir)
+		await worktreeGit.checkoutLocalBranch("main")
+
+		// await this.writeExcludeFile()
+		await this.stageAll(worktreeGit)
+		await worktreeGit.raw("--work-tree", this.worktreeDir, "commit", "-m", "Initial commit", "--allow-empty")
+		return await worktreeGit.revparse(["HEAD"])
+	}
+
+	protected override stagePath(git: SimpleGit, path: string) {
+		return git.raw("--work-tree", this.worktreeDir, "add", path)
+	}
+
+	protected override async getShadowGitConfigWorktree(git: SimpleGit) {
+		if (!this.shadowGitConfigWorktree) {
+			try {
+				const workspace = await fs.readFile(path.join(this.configDir, "workspace"), "utf-8")
+				this.shadowGitConfigWorktree = workspace || undefined
+			} catch (error) {
+				this.log(
+					`[${this.constructor.name}#getShadowGitConfigWorktree] failed to get workspace: ${error instanceof Error ? error.message : String(error)}`,
+				)
+			}
+		}
+
+		return this.shadowGitConfigWorktree
+	}
+
+	public static create({ taskId, workspaceDir, shadowDir, log = console.log }: CheckpointServiceOptions) {
+		const workspaceHash = this.hashWorkspaceDir(workspaceDir)
+		const checkpointsDir = path.join(shadowDir, "checkpoints", workspaceHash, "parent")
+		return new WorktreeCheckpointService(taskId, checkpointsDir, workspaceDir, log)
+	}
+}

src/services/checkpoints/__tests__/RepoPerWorkspaceCheckpointService.test.ts

@@ -0,0 +1,64 @@
+// npx jest src/services/checkpoints/__tests__/RepoPerWorkspaceCheckpointService.test.ts
+
+import fs from "fs/promises"
+import path from "path"
+import os from "os"
+
+import { initWorkspaceRepo } from "./initWorkspaceRepo"
+
+import { RepoPerWorkspaceCheckpointService } from "../RepoPerWorkspaceCheckpointService"
+
+jest.mock("globby", () => ({
+	globby: jest.fn().mockResolvedValue([]),
+}))
+
+describe("RepoPerWorkspaceCheckpointService", () => {
+	const tmpDir = path.join(os.tmpdir(), RepoPerWorkspaceCheckpointService.name)
+	const globalStorageDir = path.join(tmpDir, "globalStorage")
+	const workspaceDir = path.join(tmpDir, "workspace")
+	const log = () => {}
+
+	beforeEach(async () => {
+		await initWorkspaceRepo({ workspaceDir })
+		await fs.mkdir(globalStorageDir, { recursive: true })
+	})
+
+	afterEach(async () => {
+		await fs.rm(globalStorageDir, { recursive: true, force: true })
+		await fs.rm(workspaceDir, { recursive: true, force: true })
+	})
+
+	it("does not achieve isolation", async () => {
+		const service1 = RepoPerWorkspaceCheckpointService.create({
+			taskId: "task1",
+			shadowDir: globalStorageDir,
+			workspaceDir,
+			log,
+		})
+		await service1.initShadowGit()
+
+		await fs.writeFile(path.join(workspaceDir, "foo.txt"), "foo")
+		const commit1 = await service1.saveCheckpoint("foo")
+		expect(commit1?.commit).toBeTruthy()
+
+		const service2 = RepoPerWorkspaceCheckpointService.create({
+			taskId: "task2",
+			shadowDir: globalStorageDir,
+			workspaceDir,
+			log,
+		})
+		await service2.initShadowGit()
+
+		await fs.writeFile(path.join(workspaceDir, "bar.txt"), "bar")
+		const commit2 = await service2.saveCheckpoint("bar")
+		expect(commit2?.commit).toBeTruthy()
+
+		const diff = await service1.getDiff({ to: commit1!.commit })
+		expect(diff).toHaveLength(1)
+
+		expect(await fs.readFile(path.join(workspaceDir, "foo.txt"), "utf-8")).toBe("foo")
+
+		// Argh! This should not happen!
+		expect(fs.readFile(path.join(workspaceDir, "bar.txt"), "utf-8")).rejects.toThrow(/no such file or directory/)
+	})
+})

src/services/checkpoints/__tests__/ShadowCheckpointService.test.ts

@@ -9,6 +9,8 @@ import { simpleGit, SimpleGit } from "simple-git"
 
 import { fileExistsAtPath } from "../../../utils/fs"
 
+import { initWorkspaceRepo } from "./initWorkspaceRepo"
+
 import { ShadowCheckpointService } from "../ShadowCheckpointService"
 import { RepoPerTaskCheckpointService } from "../RepoPerTaskCheckpointService"
 import { RepoPerWorkspaceCheckpointService } from "../RepoPerWorkspaceCheckpointService"
@@ -17,40 +19,7 @@ jest.mock("globby", () => ({
 	globby: jest.fn().mockResolvedValue([]),
 }))
 
-const tmpDir = path.join(os.tmpdir(), "CheckpointService")
-
-const initWorkspaceRepo = async ({
-	workspaceDir,
-	userName = "Roo Code",
-	userEmail = "[email protected]",
-	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 }
-}
+export const tmpDir = path.join(os.tmpdir(), "CheckpointService")
 
 describe.each([
 	[RepoPerTaskCheckpointService, "RepoPerTaskCheckpointService"],