Skip to content

Commit 411e1ff

Browse files
ctecte
committed
Some improvements, with tests
1 parent dec44fe commit 411e1ff

File tree

2 files changed

+96
-48
lines changed

2 files changed

+96
-48
lines changed

‎src/integrations/checkpoints/LocalCheckpointer.ts‎

Lines changed: 56 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ export class LocalCheckpointer {
2626
const checkpointer = new LocalCheckpointer(options)
2727
await checkpointer.ensureGitInstalled()
2828
await checkpointer.ensureGitRepo()
29+
await checkpointer.initGitConfig()
2930
await checkpointer.initHiddenBranch()
3031
return checkpointer
3132
}
@@ -46,18 +47,6 @@ export class LocalCheckpointer {
4647
this.git = simpleGit(options)
4748
}
4849

49-
/**
50-
* Initialize git configuration. Should be called after constructor.
51-
*/
52-
private async initGitConfig(): Promise<void> {
53-
try {
54-
await this.git.addConfig("user.name", "Roo Code")
55-
await this.git.addConfig("user.email", "[email protected]")
56-
} catch (err) {
57-
throw new Error(`Failed to configure Git: ${err instanceof Error ? err.message : String(err)}`)
58-
}
59-
}
60-
6150
/**
6251
* Ensure that Git is installed.
6352
*/
@@ -85,6 +74,35 @@ export class LocalCheckpointer {
8574
}
8675
}
8776

77+
/**
78+
* Initialize git configuration. Should be called after constructor.
79+
*/
80+
private async initGitConfig() {
81+
try {
82+
await this.git.addConfig("user.name", "Roo Code")
83+
await this.git.addConfig("user.email", "[email protected]")
84+
} catch (err) {
85+
throw new Error(`Failed to configure Git: ${err instanceof Error ? err.message : String(err)}`)
86+
}
87+
}
88+
89+
/**
90+
* Create the hidden branch if it doesn't exist. Otherwise, do nothing.
91+
* If the branch is missing, we base it off the main branch.
92+
*/
93+
private async initHiddenBranch(): Promise<void> {
94+
// Check if the branch already exists.
95+
const branchSummary = await this.git.branch()
96+
97+
if (!branchSummary.all.includes(this.hiddenBranch)) {
98+
// Create the new branch from main.
99+
await this.git.checkoutBranch(this.hiddenBranch, this.mainBranch)
100+
101+
// Switch back to main.
102+
await this.git.checkout(this.mainBranch)
103+
}
104+
}
105+
88106
private async pushStash() {
89107
try {
90108
const status = await this.git.status()
@@ -145,23 +163,6 @@ export class LocalCheckpointer {
145163
}
146164
}
147165

148-
/**
149-
* Create the hidden branch if it doesn't exist. Otherwise, do nothing.
150-
* If the branch is missing, we base it off the main branch.
151-
*/
152-
private async initHiddenBranch(): Promise<void> {
153-
// Check if the branch already exists.
154-
const branchSummary = await this.git.branch()
155-
156-
if (!branchSummary.all.includes(this.hiddenBranch)) {
157-
// Create the new branch from main.
158-
await this.git.checkoutBranch(this.hiddenBranch, this.mainBranch)
159-
160-
// Switch back to main.
161-
await this.git.checkout(this.mainBranch)
162-
}
163-
}
164-
165166
/**
166167
* List commits on the hidden branch as checkpoints.
167168
* We can parse the commit log to build an array of `Checkpoint`.
@@ -177,7 +178,7 @@ export class LocalCheckpointer {
177178
}
178179

179180
/**
180-
* Commit changes in the working directory (on the hidden branch) as a new checkpoint.\
181+
* Commit changes in the working directory (on the hidden branch) as a new checkpoint.
181182
* Preserves the current state of the main branch.
182183
*/
183184
public async saveCheckpoint(message: string) {
@@ -194,10 +195,25 @@ export class LocalCheckpointer {
194195
}
195196

196197
try {
198+
// Get the latest commit on the hidden branch before we reset it.
199+
const latestHash = await this.git.revparse([this.hiddenBranch])
200+
201+
// Reset hidden branch to match main and apply the pending changes.
197202
await this.git.checkout(this.hiddenBranch)
198-
await this.git.reset(["--hard", this.mainBranch]) // Reset hidden branch to match main
199-
await this.applyStash() // Apply the stashed changes
200-
await this.git.add(["."]) // Stage everything
203+
await this.git.reset(["--hard", this.mainBranch])
204+
await this.applyStash()
205+
206+
// If there are no changes, we don't need to commit.
207+
const diff = await this.git.diff([latestHash])
208+
209+
if (!diff) {
210+
await this.git.checkout(this.mainBranch)
211+
await this.popStash()
212+
return undefined
213+
}
214+
215+
// Otherwise, commit the changes.
216+
await this.git.add(["."])
201217
const commit = await this.git.commit(message)
202218
await this.git.checkout(this.mainBranch)
203219
await this.popStash()
@@ -222,10 +238,13 @@ export class LocalCheckpointer {
222238
throw new Error(`Must be on ${this.mainBranch} branch to restore checkpoints. Currently on: ${branch}`)
223239
}
224240

225-
// Discard any pending changes. Note that these should already be preserved
226-
// as a checkpoint, but we should verify that.
227-
await this.pushStash()
228-
await this.dropStash()
241+
// Persist pending changes in a checkpoint and then discard them.
242+
const commit = await this.saveCheckpoint(`restoreCheckpoint ${commitHash}`)
243+
244+
if (commit) {
245+
await this.pushStash()
246+
await this.dropStash()
247+
}
229248

230249
await this.git.raw(["restore", "--source", commitHash, "--worktree", "--", "."])
231250
}

‎src/integrations/checkpoints/__tests__/LocalCheckpointer.test.ts‎

Lines changed: 40 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -9,40 +9,38 @@ import { CommitResult, simpleGit, SimpleGit } from "simple-git"
99
import { LocalCheckpointer } from "../LocalCheckpointer"
1010

1111
describe("LocalCheckpointer", () => {
12-
let checkpointer: LocalCheckpointer
13-
let tempDir: string
1412
let git: SimpleGit
1513
let testFile: string
1614
let initialCommit: CommitResult
17-
const mainBranch = "my_branch"
15+
const mainBranch = "main"
1816
const hiddenBranch = "checkpoints"
17+
let checkpointer: LocalCheckpointer
1918

2019
beforeEach(async () => {
2120
// Create a temporary directory for testing.
22-
tempDir = path.join(os.tmpdir(), `checkpointer-test-${Date.now()}`)
23-
await fs.mkdir(tempDir)
21+
const workspacePath = path.join(os.tmpdir(), `checkpointer-test-${Date.now()}`)
22+
await fs.mkdir(workspacePath)
2423

2524
// Initialize git repo.
26-
git = simpleGit(tempDir)
25+
git = simpleGit(workspacePath)
2726
await git.init(["--initial-branch", mainBranch])
2827
await git.addConfig("user.name", "Roo Code")
2928
await git.addConfig("user.email", "[email protected]")
3029

3130
// Create test file.
32-
testFile = path.join(tempDir, "test.txt")
31+
testFile = path.join(workspacePath, "test.txt")
3332
await fs.writeFile(testFile, "Hello, world!")
3433

3534
// Create initial commit.
3635
await git.add(".")
3736
initialCommit = await git.commit("Initial commit")!
3837

3938
// Create checkpointer instance.
40-
checkpointer = await LocalCheckpointer.create({ workspacePath: tempDir, mainBranch, hiddenBranch })
39+
checkpointer = await LocalCheckpointer.create({ workspacePath, mainBranch, hiddenBranch })
4140
})
4241

4342
afterEach(async () => {
44-
// Clean up temporary directory.
45-
await fs.rm(tempDir, { recursive: true, force: true })
43+
await fs.rm(checkpointer.workspacePath, { recursive: true, force: true })
4644
})
4745

4846
it("creates a hidden branch on initialization", async () => {
@@ -76,7 +74,6 @@ describe("LocalCheckpointer", () => {
7674
const commit2 = await checkpointer.saveCheckpoint("Second checkpoint")
7775
expect(commit2?.commit).toBeTruthy()
7876
const details2 = await git.show([commit2!.commit])
79-
console.log(details2)
8077
expect(details2).toContain("-Hello, world!")
8178
expect(details2).toContain("+Hola, world!")
8279

@@ -92,4 +89,36 @@ describe("LocalCheckpointer", () => {
9289
await checkpointer.restoreCheckpoint(initialCommit.commit)
9390
expect(await fs.readFile(testFile, "utf-8")).toBe("Hello, world!")
9491
})
92+
93+
it("does nothing if no changes are made since the last checkpoint", async () => {
94+
await fs.writeFile(testFile, "Ahoy, world!")
95+
const commit = await checkpointer.saveCheckpoint("First checkpoint")
96+
expect(commit?.commit).toBeTruthy()
97+
98+
const commit2 = await checkpointer.saveCheckpoint("Second checkpoint")
99+
expect(commit2?.commit).toBeFalsy()
100+
})
101+
102+
it("preserves pending changes when restoring a checkpoint", async () => {
103+
await fs.writeFile(testFile, "Ahoy, world!")
104+
const commit1 = await checkpointer.saveCheckpoint("First checkpoint")
105+
expect(commit1?.commit).toBeTruthy()
106+
107+
await fs.writeFile(testFile, "Hola, world!")
108+
const commit2 = await checkpointer.saveCheckpoint("Second checkpoint")
109+
expect(commit2?.commit).toBeTruthy()
110+
111+
await fs.writeFile(testFile, "Bonjour, world!")
112+
113+
// Restore first checkpoint - this should create a new checkpoint with
114+
// the pending changes.
115+
await checkpointer.restoreCheckpoint(commit1!.commit)
116+
117+
// Verify the pending changes were saved as a checkpoint.
118+
const checkpoints = await checkpointer.listCheckpoints()
119+
expect(checkpoints[0].message).toBe(`restoreCheckpoint ${commit1!.commit}`)
120+
121+
// Verify the content is now from checkpoint1.
122+
expect(await fs.readFile(testFile, "utf-8")).toBe("Ahoy, world!")
123+
})
95124
})

0 commit comments

Comments
 (0)