Skip to content

Commit 618964d

Browse files
ctecte
committed
More progress
1 parent ba0a076 commit 618964d

File tree

4 files changed

+198
-54
lines changed

4 files changed

+198
-54
lines changed

‎src/services/checkpoints/ShadowCheckpointService.ts‎

Lines changed: 59 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -20,12 +20,12 @@ export abstract class ShadowCheckpointService extends EventEmitter {
2020
public readonly checkpointsDir: string
2121
public readonly workspaceDir: string
2222

23+
protected configDir: string
24+
protected readonly log: (message: string) => void
25+
2326
protected _checkpoints: string[] = []
2427
protected _baseHash?: string
25-
26-
protected readonly dotGitDir: string
2728
protected git?: SimpleGit
28-
protected readonly log: (message: string) => void
2929
protected shadowGitConfigWorktree?: string
3030

3131
public get baseHash() {
@@ -56,8 +56,8 @@ export abstract class ShadowCheckpointService extends EventEmitter {
5656
this.taskId = taskId
5757
this.checkpointsDir = checkpointsDir
5858
this.workspaceDir = workspaceDir
59+
this.configDir = path.join(this.checkpointsDir, ".git")
5960

60-
this.dotGitDir = path.join(this.checkpointsDir, ".git")
6161
this.log = log
6262
}
6363

@@ -66,47 +66,16 @@ export abstract class ShadowCheckpointService extends EventEmitter {
6666
throw new Error("Shadow git repo already initialized")
6767
}
6868

69-
await fs.mkdir(this.checkpointsDir, { recursive: true })
70-
const git = simpleGit(this.checkpointsDir)
71-
const gitVersion = await git.version()
72-
this.log(`[${this.constructor.name}#create] git = ${gitVersion}`)
73-
74-
let created = false
7569
const startTime = Date.now()
76-
77-
if (await fileExistsAtPath(this.dotGitDir)) {
78-
this.log(`[${this.constructor.name}#initShadowGit] shadow git repo already exists at ${this.dotGitDir}`)
79-
const worktree = await this.getShadowGitConfigWorktree(git)
80-
81-
if (worktree !== this.workspaceDir) {
82-
throw new Error(
83-
`Checkpoints can only be used in the original workspace: ${worktree} !== ${this.workspaceDir}`,
84-
)
85-
}
86-
87-
await this.writeExcludeFile()
88-
this.baseHash = await git.revparse(["HEAD"])
89-
} else {
90-
this.log(`[${this.constructor.name}#initShadowGit] creating shadow git repo at ${this.checkpointsDir}`)
91-
await git.init()
92-
await git.addConfig("core.worktree", this.workspaceDir) // Sets the working tree to the current workspace.
93-
await git.addConfig("commit.gpgSign", "false") // Disable commit signing for shadow repo.
94-
await git.addConfig("user.name", "Roo Code")
95-
await git.addConfig("user.email", "[email protected]")
96-
await this.writeExcludeFile()
97-
await this.stageAll(git)
98-
const { commit } = await git.commit("initial commit", { "--allow-empty": null })
99-
this.baseHash = commit
100-
created = true
101-
}
102-
70+
const { git, baseHash, created } = await this.initializeShadowRepo()
10371
const duration = Date.now() - startTime
10472

10573
this.log(
106-
`[${this.constructor.name}#initShadowGit] initialized shadow repo with base commit ${this.baseHash} in ${duration}ms`,
74+
`[${this.constructor.name}#initShadowGit] initialized shadow repo with base commit ${baseHash} in ${duration}ms`,
10775
)
10876

10977
this.git = git
78+
this.baseHash = baseHash
11079

11180
await onInit?.()
11281

@@ -121,22 +90,68 @@ export abstract class ShadowCheckpointService extends EventEmitter {
12190
return { created, duration }
12291
}
12392

93+
protected isShadowRepoAvailable() {
94+
return fileExistsAtPath(this.configDir)
95+
}
96+
97+
protected async initializeShadowRepo() {
98+
await fs.mkdir(this.checkpointsDir, { recursive: true })
99+
const git = simpleGit(this.checkpointsDir)
100+
const gitVersion = await git.version()
101+
this.log(`[${this.constructor.name}#initializeShadowRepo] git = ${gitVersion}`)
102+
const exists = await this.isShadowRepoAvailable()
103+
console.log(`[${this.constructor.name}#initializeShadowRepo] exists = ${exists} [${this.configDir}]`)
104+
const baseHash = exists ? await this.checkShadowRepo(git) : await this.createShadowRepo(git)
105+
return { git, baseHash, created: !exists }
106+
}
107+
108+
protected async checkShadowRepo(git: SimpleGit) {
109+
this.log(`[${this.constructor.name}#checkShadowRepo] checking existing shadow repo at ${this.configDir}`)
110+
const worktree = await this.getShadowGitConfigWorktree(git)
111+
112+
if (worktree !== this.workspaceDir) {
113+
throw new Error(
114+
`Checkpoints can only be used in the original workspace: ${worktree} !== ${this.workspaceDir}`,
115+
)
116+
}
117+
118+
await this.writeExcludeFile()
119+
return await git.revparse(["HEAD"])
120+
}
121+
122+
protected async createShadowRepo(git: SimpleGit) {
123+
this.log(`[${this.constructor.name}#createShadowRepo] creating new shadow repo at ${this.checkpointsDir}`)
124+
await git.init()
125+
await git.addConfig("core.worktree", this.workspaceDir) // Sets the working tree to the current workspace.
126+
await git.addConfig("commit.gpgSign", "false") // Disable commit signing for shadow repo.
127+
await git.addConfig("user.name", "Roo Code")
128+
await git.addConfig("user.email", "[email protected]")
129+
await this.writeExcludeFile()
130+
await this.stageAll(git)
131+
const result = await git.commit("initial commit", { "--allow-empty": null })
132+
return result.commit
133+
}
134+
124135
// Add basic excludes directly in git config, while respecting any
125136
// .gitignore in the workspace.
126137
// .git/info/exclude is local to the shadow git repo, so it's not
127138
// shared with the main repo - and won't conflict with user's
128139
// .gitignore.
129140
protected async writeExcludeFile() {
130-
await fs.mkdir(path.join(this.dotGitDir, "info"), { recursive: true })
141+
await fs.mkdir(path.join(this.configDir, "info"), { recursive: true })
131142
const patterns = await getExcludePatterns(this.workspaceDir)
132-
await fs.writeFile(path.join(this.dotGitDir, "info", "exclude"), patterns.join("\n"))
143+
await fs.writeFile(path.join(this.configDir, "info", "exclude"), patterns.join("\n"))
144+
}
145+
146+
protected stagePath(git: SimpleGit, path: string) {
147+
return git.add(path)
133148
}
134149

135-
private async stageAll(git: SimpleGit) {
150+
protected async stageAll(git: SimpleGit) {
136151
await this.renameNestedGitRepos(true)
137152

138153
try {
139-
await git.add(".")
154+
await this.stagePath(git, ".")
140155
} catch (error) {
141156
this.log(
142157
`[${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 {
149164
// Since we use git to track checkpoints, we need to temporarily disable
150165
// nested git repos to work around git's requirement of using submodules for
151166
// nested repos.
152-
private async renameNestedGitRepos(disable: boolean) {
167+
protected async renameNestedGitRepos(disable: boolean) {
153168
// Find all .git directories that are not at the root level.
154169
const gitPaths = await globby("**/.git" + (disable ? "" : GIT_DISABLED_SUFFIX), {
155170
cwd: this.workspaceDir,
@@ -186,7 +201,7 @@ export abstract class ShadowCheckpointService extends EventEmitter {
186201
}
187202
}
188203

189-
private async getShadowGitConfigWorktree(git: SimpleGit) {
204+
protected async getShadowGitConfigWorktree(git: SimpleGit) {
190205
if (!this.shadowGitConfigWorktree) {
191206
try {
192207
this.shadowGitConfigWorktree = (await git.getConfig("core.worktree")).value || undefined

‎src/services/checkpoints/WorktreeCheckpointService.ts‎

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,75 @@
1+
import fs from "fs/promises"
2+
import * as path from "path"
3+
4+
import simpleGit, { SimpleGit } from "simple-git"
5+
6+
import { fileExistsAtPath } from "../../utils/fs"
7+
8+
import { CheckpointServiceOptions } from "./types"
9+
import { ShadowCheckpointService } from "./ShadowCheckpointService"
10+
11+
export class WorktreeCheckpointService extends ShadowCheckpointService {
12+
private readonly worktreeDir: string
13+
14+
constructor(taskId: string, checkpointsDir: string, workspaceDir: string, log = console.log) {
15+
super(taskId, checkpointsDir, workspaceDir, log)
16+
this.configDir = this.checkpointsDir
17+
this.worktreeDir = path.join(path.dirname(this.checkpointsDir), taskId)
18+
}
19+
20+
protected override isShadowRepoAvailable() {
21+
return fileExistsAtPath(path.join(this.configDir, "workspace"))
22+
}
23+
24+
protected override async initializeShadowRepo() {
25+
const result = await super.initializeShadowRepo()
26+
const git = simpleGit(this.worktreeDir)
27+
return { ...result, git }
28+
}
29+
30+
protected override async createShadowRepo(git: SimpleGit) {
31+
this.log(`[${this.constructor.name}#createShadowRepo] creating bare shadow git repo at ${this.configDir}`)
32+
await git.init(["--bare"])
33+
34+
await fs.writeFile(path.join(this.configDir, "workspace"), this.workspaceDir)
35+
36+
await git.addConfig("commit.gpgSign", "false") // Disable commit signing for shadow repo.
37+
await git.addConfig("user.name", "Roo Code")
38+
await git.addConfig("user.email", "[email protected]")
39+
40+
await git.raw(["worktree", "add", "--orphan", this.worktreeDir])
41+
42+
const worktreeGit = simpleGit(this.worktreeDir)
43+
await worktreeGit.checkoutLocalBranch("main")
44+
45+
// await this.writeExcludeFile()
46+
await this.stageAll(worktreeGit)
47+
await worktreeGit.raw("--work-tree", this.worktreeDir, "commit", "-m", "Initial commit", "--allow-empty")
48+
return await worktreeGit.revparse(["HEAD"])
49+
}
50+
51+
protected override stagePath(git: SimpleGit, path: string) {
52+
return git.raw("--work-tree", this.worktreeDir, "add", path)
53+
}
54+
55+
protected override async getShadowGitConfigWorktree(git: SimpleGit) {
56+
if (!this.shadowGitConfigWorktree) {
57+
try {
58+
const workspace = await fs.readFile(path.join(this.configDir, "workspace"), "utf-8")
59+
this.shadowGitConfigWorktree = workspace || undefined
60+
} catch (error) {
61+
this.log(
62+
`[${this.constructor.name}#getShadowGitConfigWorktree] failed to get workspace: ${error instanceof Error ? error.message : String(error)}`,
63+
)
64+
}
65+
}
66+
67+
return this.shadowGitConfigWorktree
68+
}
69+
70+
public static create({ taskId, workspaceDir, shadowDir, log = console.log }: CheckpointServiceOptions) {
71+
const workspaceHash = this.hashWorkspaceDir(workspaceDir)
72+
const checkpointsDir = path.join(shadowDir, "checkpoints", workspaceHash, "parent")
73+
return new WorktreeCheckpointService(taskId, checkpointsDir, workspaceDir, log)
74+
}
75+
}

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

Lines changed: 18 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -13,32 +13,40 @@ jest.mock("globby", () => ({
1313
}))
1414

1515
describe("RepoPerWorkspaceCheckpointService", () => {
16-
const tmpDir = path.join(os.tmpdir(), "RepoPerWorkspaceCheckpointService")
17-
const shadowDir = path.join(tmpDir, "shadow-dir")
18-
const workspaceDir = path.join(tmpDir, "workspace-dir")
19-
const log = console.log
16+
const tmpDir = path.join(os.tmpdir(), RepoPerWorkspaceCheckpointService.name)
17+
const globalStorageDir = path.join(tmpDir, "globalStorage")
18+
const workspaceDir = path.join(tmpDir, "workspace")
19+
const log = () => {}
2020

2121
beforeEach(async () => {
2222
await initWorkspaceRepo({ workspaceDir })
23-
await fs.mkdir(shadowDir, { recursive: true })
23+
await fs.mkdir(globalStorageDir, { recursive: true })
2424
})
2525

2626
afterEach(async () => {
27-
await fs.rm(shadowDir, { recursive: true, force: true })
27+
await fs.rm(globalStorageDir, { recursive: true, force: true })
2828
await fs.rm(workspaceDir, { recursive: true, force: true })
2929
})
3030

3131
it("does not achieve isolation", async () => {
32-
const task1 = "task1"
33-
const service1 = RepoPerWorkspaceCheckpointService.create({ taskId: task1, shadowDir, workspaceDir, log })
32+
const service1 = RepoPerWorkspaceCheckpointService.create({
33+
taskId: "task1",
34+
shadowDir: globalStorageDir,
35+
workspaceDir,
36+
log,
37+
})
3438
await service1.initShadowGit()
3539

3640
await fs.writeFile(path.join(workspaceDir, "foo.txt"), "foo")
3741
const commit1 = await service1.saveCheckpoint("foo")
3842
expect(commit1?.commit).toBeTruthy()
3943

40-
const task2 = "task2"
41-
const service2 = RepoPerWorkspaceCheckpointService.create({ taskId: task2, shadowDir, workspaceDir, log })
44+
const service2 = RepoPerWorkspaceCheckpointService.create({
45+
taskId: "task2",
46+
shadowDir: globalStorageDir,
47+
workspaceDir,
48+
log,
49+
})
4250
await service2.initShadowGit()
4351

4452
await fs.writeFile(path.join(workspaceDir, "bar.txt"), "bar")

‎src/services/checkpoints/__tests__/WorktreeCheckpointService.test.ts‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
// npx jest src/services/checkpoints/__tests__/WorktreeCheckpointService.test.ts
2+
3+
import fs from "fs/promises"
4+
import path from "path"
5+
import os from "os"
6+
7+
import { initWorkspaceRepo } from "./initWorkspaceRepo"
8+
9+
import { WorktreeCheckpointService } from "../WorktreeCheckpointService"
10+
11+
jest.mock("globby", () => ({
12+
globby: jest.fn().mockResolvedValue([]),
13+
}))
14+
15+
describe("WorktreeCheckpointService", () => {
16+
const tmpDir = path.join(os.tmpdir(), WorktreeCheckpointService.name)
17+
const globalStorageDir = path.join(tmpDir, "globalStorage")
18+
const workspaceDir = path.join(tmpDir, "workspace")
19+
const log = console.log
20+
21+
const cleanup = async () => {
22+
await fs.rm(globalStorageDir, { recursive: true, force: true })
23+
await fs.rm(workspaceDir, { recursive: true, force: true })
24+
}
25+
26+
beforeEach(async () => {
27+
await cleanup()
28+
await initWorkspaceRepo({ workspaceDir })
29+
await fs.mkdir(globalStorageDir, { recursive: true })
30+
})
31+
32+
afterAll(async () => {
33+
await cleanup()
34+
})
35+
36+
it("achieves isolation", async () => {
37+
const service1 = WorktreeCheckpointService.create({
38+
taskId: "task1",
39+
shadowDir: globalStorageDir,
40+
workspaceDir,
41+
log,
42+
})
43+
await service1.initShadowGit()
44+
console.log(service1.checkpointsDir)
45+
})
46+
})

0 commit comments

Comments
 (0)