Skip to content

Commit 9d83606

Browse files
ctecte
committed
New checkpoints algorithm
1 parent ca54d1e commit 9d83606

File tree

2 files changed

+186
-39
lines changed

2 files changed

+186
-39
lines changed

‎src/services/checkpoints/CheckpointService.ts‎

Lines changed: 182 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,12 @@ export type CheckpointServiceOptions = {
1111
log?: (message: string) => void
1212
}
1313

14+
export type StagedFile = {
15+
path: string
16+
isStaged: boolean
17+
isPartiallyStaged: boolean
18+
}
19+
1420
/**
1521
* The CheckpointService provides a mechanism for storing a snapshot of the
1622
* current VSCode workspace each time a Roo Code tool is executed. It uses Git
@@ -23,11 +29,11 @@ export type CheckpointServiceOptions = {
2329
* - A hidden branch for storing checkpoints.
2430
*
2531
* Saving a checkpoint:
26-
* - Current changes are stashed (including untracked files).
32+
* - A temporary branch is created to store the current state.
33+
* - All changes (including untracked files) are staged and committed on the temp branch.
2734
* - The hidden branch is reset to match main.
28-
* - Stashed changes are applied and committed as a checkpoint on the hidden
29-
* branch.
30-
* - We return to the main branch with the original state restored.
35+
* - The temporary branch commit is cherry-picked onto the hidden branch.
36+
* - The workspace is restored to its original state and the temp branch is deleted.
3137
*
3238
* Restoring a checkpoint:
3339
* - The workspace is restored to the state of the specified checkpoint using
@@ -37,6 +43,7 @@ export type CheckpointServiceOptions = {
3743
* - Non-destructive version control (main branch remains untouched).
3844
* - Preservation of the full history of checkpoints.
3945
* - Safe restoration to any previous checkpoint.
46+
* - Atomic checkpoint operations with proper error recovery.
4047
*
4148
* NOTES
4249
*
@@ -122,51 +129,188 @@ export class CheckpointService {
122129
return result
123130
}
124131

132+
private async restoreMain({
133+
branch,
134+
stashSha,
135+
force = false,
136+
}: {
137+
branch: string
138+
stashSha: string
139+
force?: boolean
140+
}) {
141+
if (force) {
142+
await this.git.checkout(["-f", this.mainBranch])
143+
} else {
144+
await this.git.checkout(this.mainBranch)
145+
}
146+
147+
if (stashSha) {
148+
console.log(`[restoreMain] applying stash ${stashSha}`)
149+
await this.git.raw(["stash", "apply", "--index", stashSha])
150+
}
151+
152+
console.log(`[restoreMain] restoring from ${branch}`)
153+
await this.git.raw(["restore", "--source", branch, "--worktree", "--", "."])
154+
}
155+
125156
public async saveCheckpoint(message: string) {
126157
await this.ensureBranch(this.mainBranch)
127158

128-
// Create temporary branch with all current changes
129-
const tempBranch = `${CheckpointService.STASH_BRANCH}-${Date.now()}`
130-
await this.git.checkout(["-b", tempBranch])
131-
159+
const status = await this.git.status()
160+
console.log(`[saveCheckpoint] status: ${JSON.stringify(status)}`)
161+
const stashSha = (await this.git.raw(["stash", "create"])).trim()
162+
console.log(`[saveCheckpoint] stashSha: ${stashSha}`)
163+
const latestHash = await this.git.revparse([this.hiddenBranch])
164+
165+
/**
166+
* PHASE: Create stash
167+
* Mutations:
168+
* - Create branch
169+
* - Change branch
170+
*/
171+
const stashBranch = `${CheckpointService.STASH_BRANCH}-${Date.now()}`
172+
await this.git.checkout(["-b", stashBranch])
173+
this.log(`[saveCheckpoint] created and checked out ${stashBranch}`)
174+
175+
/**
176+
* Phase: Stage stash
177+
* Mutations: None
178+
* Recovery:
179+
* - UNDO: Create branch
180+
* - UNDO: Change branch
181+
*/
132182
try {
133-
// Stage and commit all changes to temporary branch
134183
await this.git.add(["-A"])
135-
await this.git.commit(message, undefined, {
136-
"--allow-empty": null,
137-
"--no-verify": null,
138-
})
184+
const status = await this.git.status()
185+
this.log(`[saveCheckpoint] status: ${JSON.stringify(status)}`)
186+
} catch (err) {
187+
await this.git.checkout(["-f", this.mainBranch])
188+
await this.git.branch(["-D", stashBranch]).catch(() => {})
139189

140-
// Get the latest commit on the hidden branch before we reset it
141-
const latestHash = await this.git.revparse([this.hiddenBranch])
142-
await this.git.checkout(this.hiddenBranch)
190+
throw new Error(
191+
`[saveCheckpoint] Failed in stage stash phase: ${err instanceof Error ? err.message : String(err)}`,
192+
)
193+
}
143194

144-
try {
145-
// Reset hidden branch to match main and apply the changes
146-
await this.git.reset(["--hard", this.mainBranch])
195+
/**
196+
* Phase: Commit stash
197+
* Mutations:
198+
* - Commit stash
199+
* - Change branch
200+
* Recovery:
201+
* - UNDO: Create branch
202+
* - UNDO: Change branch
203+
*/
204+
try {
205+
// TODO: Add a test to see if empty commits break this.
206+
const tempCommit = await this.git.commit(message, undefined, { "--no-verify": null })
207+
this.log(`[saveCheckpoint] tempCommit: ${message} -> ${JSON.stringify(tempCommit)}`)
208+
} catch (err) {
209+
await this.git.checkout(["-f", this.mainBranch])
210+
await this.git.branch(["-D", stashBranch]).catch(() => {})
147211

148-
// Cherry-pick the temporary commit
149-
const commit = await this.git.raw(["cherry-pick", tempBranch])
212+
throw new Error(
213+
`[saveCheckpoint] Failed in stash commit phase: ${err instanceof Error ? err.message : String(err)}`,
214+
)
215+
}
150216

151-
// Return to main branch and cleanup
152-
await this.git.checkout(this.mainBranch)
153-
await this.git.branch(["-D", tempBranch])
217+
/**
218+
* PHASE: Diff
219+
* Mutations:
220+
* - Checkout hidden branch
221+
* Recovery:
222+
* - UNDO: Create branch
223+
* - UNDO: Change branch
224+
* - UNDO: Commit stash
225+
*/
226+
let diff
154227

155-
this.currentCheckpoint = commit
156-
return { commit }
228+
try {
229+
diff = await this.git.diff([latestHash, stashBranch])
230+
} catch (err) {
231+
await this.restoreMain({ branch: stashBranch, stashSha, force: true })
232+
await this.git.branch(["-D", stashBranch]).catch(() => {})
233+
234+
throw new Error(
235+
`[saveCheckpoint] Failed in diff phase: ${err instanceof Error ? err.message : String(err)}`,
236+
)
237+
}
238+
239+
if (!diff) {
240+
this.log("[saveCheckpoint] no diff")
241+
await this.restoreMain({ branch: stashBranch, stashSha })
242+
await this.git.branch(["-D", stashBranch])
243+
return undefined
244+
}
245+
246+
/**
247+
* PHASE: Reset
248+
* Mutations:
249+
* - Reset hidden branch
250+
* Recovery:
251+
* - UNDO: Create branch
252+
* - UNDO: Change branch
253+
* - UNDO: Commit stash
254+
*/
255+
try {
256+
await this.git.checkout(this.hiddenBranch)
257+
this.log(`[saveCheckpoint] checked out ${this.hiddenBranch}`)
258+
await this.git.reset(["--hard", this.mainBranch])
259+
this.log(`[saveCheckpoint] reset ${this.hiddenBranch}`)
260+
} catch (err) {
261+
await this.restoreMain({ branch: stashBranch, stashSha, force: true })
262+
await this.git.branch(["-D", stashBranch]).catch(() => {})
263+
264+
throw new Error(
265+
`[saveCheckpoint] Failed in reset phase: ${err instanceof Error ? err.message : String(err)}`,
266+
)
267+
}
268+
269+
/**
270+
* PHASE: Cherry pick
271+
* Mutations:
272+
* - Hidden commit (NOTE: reset on hidden branch no longer needed in
273+
* success scenario.)
274+
* Recovery:
275+
* - UNDO: Create branch
276+
* - UNDO: Change branch
277+
* - UNDO: Commit stash
278+
* - UNDO: Reset hidden branch
279+
*/
280+
let commit = ""
281+
282+
try {
283+
try {
284+
await this.git.raw(["cherry-pick", stashBranch])
157285
} catch (err) {
158-
// If something went wrong after switching to hidden branch
159-
await this.git.reset(["--hard", latestHash])
160-
await this.git.checkout(["-f", this.mainBranch])
161-
await this.git.branch(["-D", tempBranch])
162-
throw err
286+
// Check if we're in the middle of a cherry-pick.
287+
// If the cherry-pick resulted in an empty commit (e.g., only
288+
// deletions) then complete it with --allow-empty.
289+
// Otherwise, rethrow the error.
290+
if (existsSync(path.join(this.baseDir, ".git/CHERRY_PICK_HEAD"))) {
291+
await this.git.raw(["commit", "--allow-empty", "--no-edit"])
292+
} else {
293+
throw err
294+
}
163295
}
296+
297+
commit = await this.git.revparse(["HEAD"])
298+
this.currentCheckpoint = commit
299+
this.log(`[saveCheckpoint] cherry-pick commit = ${commit}`)
164300
} catch (err) {
165-
// If something went wrong before switching to hidden branch
166-
await this.git.checkout(["-f", this.mainBranch])
167-
await this.git.branch(["-D", tempBranch])
168-
throw new Error(`Failed to save checkpoint: ${err instanceof Error ? err.message : String(err)}`)
301+
await this.git.reset(["--hard", latestHash]).catch(() => {})
302+
await this.restoreMain({ branch: stashBranch, stashSha, force: true })
303+
await this.git.branch(["-D", stashBranch]).catch(() => {})
304+
305+
throw new Error(
306+
`[saveCheckpoint] Failed in cherry pick phase: ${err instanceof Error ? err.message : String(err)}`,
307+
)
169308
}
309+
310+
await this.restoreMain({ branch: stashBranch, stashSha })
311+
await this.git.branch(["-D", stashBranch])
312+
313+
return { commit }
170314
}
171315

172316
public async restoreCheckpoint(commitHash: string) {
@@ -264,12 +408,12 @@ export class CheckpointService {
264408
const currentBranch = await git.revparse(["--abbrev-ref", "HEAD"])
265409
const currentSha = await git.revparse(["HEAD"])
266410

267-
const hiddenBranch = `roo-code-checkpoints-${taskId}`
411+
const hiddenBranch = `${CheckpointService.CHECKPOINT_BRANCH}-${taskId}`
268412
const branchSummary = await git.branch()
269413

270414
if (!branchSummary.all.includes(hiddenBranch)) {
271-
await git.checkoutBranch(hiddenBranch, currentBranch) // git checkout -b <hiddenBranch> <currentBranch>
272-
await git.checkout(currentBranch) // git checkout <currentBranch>
415+
await git.checkoutBranch(hiddenBranch, currentBranch)
416+
await git.checkout(currentBranch)
273417
}
274418

275419
return { currentBranch, currentSha, hiddenBranch }

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,7 @@ describe("CheckpointService", () => {
6868

6969
git = repo.git
7070
testFile = repo.testFile
71-
service = await CheckpointService.create({ taskId, git, baseDir, log: () => {} })
71+
service = await CheckpointService.create({ taskId, git, baseDir })
7272
})
7373

7474
afterEach(async () => {
@@ -295,8 +295,11 @@ describe("CheckpointService", () => {
295295
await fs.writeFile(testFile, "I am tracked!")
296296
const untrackedFile = path.join(service.baseDir, "new.txt")
297297
await fs.writeFile(untrackedFile, "I am untracked!")
298+
298299
const commit1 = await service.saveCheckpoint("First checkpoint")
299300
expect(commit1?.commit).toBeTruthy()
301+
expect(await fs.readFile(testFile, "utf-8")).toBe("I am tracked!")
302+
expect(await fs.readFile(untrackedFile, "utf-8")).toBe("I am untracked!")
300303

301304
await fs.unlink(testFile)
302305
await fs.unlink(untrackedFile)

0 commit comments

Comments
 (0)