Skip to content

New checkpoints algorithm #939

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
mrubens merged 5 commits into from
Feb 11, 2025
Merged

New checkpoints algorithm #939

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
ca54d1e
Stashless checkpoints
cte Feb 11, 2025
9d83606
New checkpoints algorithm
cte Feb 11, 2025
971d29b
Revert debugging
cte Feb 11, 2025
a985494
Remove unused type
cte Feb 11, 2025
6dbb596
Revert debugging
cte Feb 11, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
272 changes: 174 additions & 98 deletions src/services/checkpoints/CheckpointService.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,12 @@ export type CheckpointServiceOptions = {
log?: (message: string) => void
}

export type StagedFile = {
path: string
isStaged: boolean
isPartiallyStaged: boolean
}

/**
* 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
Expand All @@ -23,11 +29,11 @@ export type CheckpointServiceOptions = {
* - A hidden branch for storing checkpoints.
*
* Saving a checkpoint:
* - Current changes are stashed (including untracked files).
* - 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.
* - Stashed changes are applied and committed as a checkpoint on the hidden
* branch.
* - We return to the main branch with the original state restored.
* - 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
Expand All @@ -37,6 +43,7 @@ export type CheckpointServiceOptions = {
* - 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
*
Expand All @@ -51,6 +58,8 @@ export type CheckpointServiceOptions = {
export class CheckpointService {
private static readonly USER_NAME = "Roo Code"
private static readonly USER_EMAIL = "[email protected]"
private static readonly CHECKPOINT_BRANCH = "roo-code-checkpoints"
private static readonly STASH_BRANCH = "roo-code-stash"

private _currentCheckpoint?: string

Expand All @@ -72,39 +81,6 @@ export class CheckpointService {
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"])

Expand Down Expand Up @@ -153,88 +129,188 @@ export class CheckpointService {
return result
}

private async restoreMain({
branch,
stashSha,
force = false,
}: {
branch: string
stashSha: string
force?: boolean
}) {
if (force) {
await this.git.checkout(["-f", this.mainBranch])
} else {
await this.git.checkout(this.mainBranch)
}

if (stashSha) {
console.log(`[restoreMain] applying stash ${stashSha}`)
Copy link
Contributor

@ellipsis-dev ellipsis-dev bot Feb 11, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Use the injected logging function (this.log) instead of console.log for consistency & centralized logging (see our Dev Standards).

Suggested change
console.log(`[restoreMain] applying stash ${stashSha}`)
this.log(`[restoreMain] applying stash ${stashSha}`)

await this.git.raw(["stash", "apply", "--index", stashSha])
}

console.log(`[restoreMain] restoring from ${branch}`)
await this.git.raw(["restore", "--source", branch, "--worktree", "--", "."])
}

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 status = await this.git.status()
console.log(`[saveCheckpoint] status: ${JSON.stringify(status)}`)
const stashSha = (await this.git.raw(["stash", "create"])).trim()
console.log(`[saveCheckpoint] stashSha: ${stashSha}`)
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])
/**
* PHASE: Create stash
* Mutations:
* - Create branch
* - Change branch
*/
const stashBranch = `${CheckpointService.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"])
const status = await this.git.status()
this.log(`[saveCheckpoint] status: ${JSON.stringify(status)}`)
} catch (err) {
await this.git.checkout(["-f", this.mainBranch])
await this.git.branch(["-D", stashBranch]).catch(() => {})

if (!diff) {
this.log(`[saveCheckpoint] No changes detected, giving up`)
return undefined
}
throw new Error(
`[saveCheckpoint] Failed in stage stash phase: ${err instanceof Error ? err.message : String(err)}`,
)
}

await this.git.checkout(this.hiddenBranch)
/**
* Phase: Commit stash
* Mutations:
* - Commit stash
* - Change branch
* Recovery:
* - UNDO: Create branch
* - UNDO: Change branch
*/
try {
// TODO: Add a test to see if empty commits break this.
const tempCommit = await this.git.commit(message, undefined, { "--no-verify": null })
this.log(`[saveCheckpoint] tempCommit: ${message} -> ${JSON.stringify(tempCommit)}`)
} catch (err) {
await this.git.checkout(["-f", this.mainBranch])
await this.git.branch(["-D", stashBranch]).catch(() => {})

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()
throw new Error(
`[saveCheckpoint] Failed in stash commit phase: ${err instanceof Error ? err.message : String(err)}`,
)
}

/**
* PHASE: Diff
* Mutations:
* - Checkout hidden branch
* Recovery:
* - UNDO: Create branch
* - UNDO: Change branch
* - UNDO: Commit stash
*/
let diff

try {
// Reset hidden branch to match main and apply the pending changes.
await this.git.reset(["--hard", this.mainBranch])
diff = await this.git.diff([latestHash, stashBranch])
} catch (err) {
await this.restoreMain({ branch: stashBranch, stashSha, force: true })
await this.git.branch(["-D", stashBranch]).catch(() => {})

if (pendingChanges) {
await this.applyStash()
}
throw new Error(
`[saveCheckpoint] Failed in diff phase: ${err instanceof Error ? err.message : String(err)}`,
)
}

// 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 diff")
await this.restoreMain({ branch: stashBranch, stashSha })
await this.git.branch(["-D", stashBranch])
return undefined
}

if (!diff) {
this.log(`[saveCheckpoint] No changes detected, resetting and giving up`)
await reset()
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) {
await this.restoreMain({ branch: stashBranch, stashSha, force: true })
await this.git.branch(["-D", stashBranch]).catch(() => {})

// 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,
})
throw new Error(
`[saveCheckpoint] Failed in reset phase: ${err instanceof Error ? err.message : String(err)}`,
)
}

await this.git.checkout(this.mainBranch)
/**
* 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 = ""

if (pendingChanges) {
await this.popStash()
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.baseDir, ".git/CHERRY_PICK_HEAD"))) {
await this.git.raw(["commit", "--allow-empty", "--no-edit"])
} else {
throw err
}
}

this.currentCheckpoint = commit.commit

return commit
commit = await this.git.revparse(["HEAD"])
this.currentCheckpoint = commit
this.log(`[saveCheckpoint] cherry-pick commit = ${commit}`)
} catch (err) {
this.log(`[saveCheckpoint] Failed to save checkpoint: ${err instanceof Error ? err.message : String(err)}`)
await this.git.reset(["--hard", latestHash]).catch(() => {})
await this.restoreMain({ branch: stashBranch, stashSha, force: true })
await this.git.branch(["-D", stashBranch]).catch(() => {})

// 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"])
throw new Error(
`[saveCheckpoint] Failed in cherry pick phase: ${err instanceof Error ? err.message : String(err)}`,
)
}

if (currentBranch.trim() !== this.mainBranch) {
await reset()
}
await this.restoreMain({ branch: stashBranch, stashSha })
await this.git.branch(["-D", stashBranch])

throw err
}
return { commit }
}

public async restoreCheckpoint(commitHash: string) {
Expand Down Expand Up @@ -332,12 +408,12 @@ export class CheckpointService {
const currentBranch = await git.revparse(["--abbrev-ref", "HEAD"])
const currentSha = await git.revparse(["HEAD"])

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

if (!branchSummary.all.includes(hiddenBranch)) {
await git.checkoutBranch(hiddenBranch, currentBranch) // git checkout -b <hiddenBranch> <currentBranch>
await git.checkout(currentBranch) // git checkout <currentBranch>
await git.checkoutBranch(hiddenBranch, currentBranch)
await git.checkout(currentBranch)
}

return { currentBranch, currentSha, hiddenBranch }
Expand Down
5 changes: 4 additions & 1 deletion src/services/checkpoints/__tests__/CheckpointService.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -68,7 +68,7 @@ describe("CheckpointService", () => {

git = repo.git
testFile = repo.testFile
service = await CheckpointService.create({ taskId, git, baseDir, log: () => {} })
service = await CheckpointService.create({ taskId, git, baseDir })
})

afterEach(async () => {
Expand Down Expand Up @@ -295,8 +295,11 @@ describe("CheckpointService", () => {
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()
expect(await fs.readFile(testFile, "utf-8")).toBe("I am tracked!")
expect(await fs.readFile(untrackedFile, "utf-8")).toBe("I am untracked!")

await fs.unlink(testFile)
await fs.unlink(untrackedFile)
Expand Down
Loading