scriptoriumadmin/Roo-Code
1
0
Fork 0
mirror of https://github.com/RooVetGit/Roo-Code.git synced 2026-08-28 05:27:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Async checkpoints

Browse source
This commit is contained in:
cte 2025-03-03 12:28:16 -08:00
parent 464d440560
commit 9e471444de
10 changed files with 392 additions and 962 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

142
src/core/Cline.ts
View file

@ -10,10 +10,10 @@ import getFolderSize from "get-folder-size"
import * as path from "path"
import { serializeError } from "serialize-error"
import * as vscode from "vscode"
import { ApiHandler, SingleCompletionHandler, buildApiHandler } from "../api"
import { ApiHandler, buildApiHandler } from "../api"
import { ApiStream } from "../api/transform/stream"
import { DIFF_VIEW_URI_SCHEME, DiffViewProvider } from "../integrations/editor/DiffViewProvider"
import { CheckpointService, CheckpointServiceFactory } from "../services/checkpoints"
import { ShadowCheckpointService } from "../services/checkpoints/ShadowCheckpointService"
import { findToolName, formatContentBlockToMarkdown } from "../integrations/misc/export-markdown"
import {
extractTextFromFile,
@ -116,7 +116,7 @@ export class Cline {
// checkpoints
enableCheckpoints: boolean = false
private checkpointService?: CheckpointService
private checkpointService?: ShadowCheckpointService
// streaming
isWaitingForFirstChunk = false
@ -747,8 +747,11 @@ export class Cline {
}
private async initiateTaskLoop(userContent: UserContent): Promise<void> {
this.initializeCheckpoints()
let nextUserContent = userContent
let includeFileDetails = true
while (!this.abort) {
const didEndLoop = await this.recursivelyMakeClineRequests(nextUserContent, includeFileDetails)
includeFileDetails = false // we only need file details the first time
@ -2773,7 +2776,7 @@ export class Cline {
}
if (isCheckpointPossible) {
await this.checkpointSave({ isFirst: false })
this.checkpointSave()
}
/*
@ -2839,13 +2842,6 @@ export class Cline {
// get previous api req's index to check token usage and determine if we need to truncate conversation history
const previousApiReqIndex = findLastIndex(this.clineMessages, (m) => m.say === "api_req_started")
// Save checkpoint if this is the first API request.
const isFirstRequest = this.clineMessages.filter((m) => m.say === "api_req_started").length === 0
if (isFirstRequest) {
await this.checkpointSave({ isFirst: true })
}
// getting verbose details is an expensive operation, it uses globby to top-down build file structure of project which for large projects can take a few seconds
// for the best UX we show a placeholder api_req_started message with a loading spinner as this happens
await this.say(
@ -3356,37 +3352,88 @@ export class Cline {
// Checkpoints
private async getCheckpointService() {
private async initializeCheckpoints() {
if (!this.enableCheckpoints) {
throw new Error("Checkpoints are disabled")
return
}
if (!this.checkpointService) {
const log = (message: string) => {
console.log(message)
try {
this.providerRef.deref()?.log(message)
} catch (err) {
// NO-OP
}
}
try {
if (this.checkpointService) {
log("[Cline#initializeCheckpoints] checkpointService already initialized")
return
}
const workspaceDir = vscode.workspace.workspaceFolders?.map((folder) => folder.uri.fsPath).at(0)
const shadowDir = this.providerRef.deref()?.context.globalStorageUri.fsPath
if (!workspaceDir) {
this.providerRef.deref()?.log("[getCheckpointService] workspace folder not found")
throw new Error("Workspace directory not found")
log("[Cline#initializeCheckpoints] workspace folder not found, disabling checkpoints")
this.enableCheckpoints = false
return
}
const shadowDir = this.providerRef.deref()?.context.globalStorageUri.fsPath
if (!shadowDir) {
this.providerRef.deref()?.log("[getCheckpointService] shadowDir not found")
throw new Error("Global storage directory not found")
log("[Cline#initializeCheckpoints] shadowDir not found, disabling checkpoints")
this.enableCheckpoints = false
return
}
this.checkpointService = await CheckpointServiceFactory.create({
strategy: "shadow",
options: {
taskId: this.taskId,
workspaceDir,
shadowDir,
log: (message) => this.providerRef.deref()?.log(message),
},
})
}
const service = await ShadowCheckpointService.create({ taskId: this.taskId, workspaceDir, shadowDir, log })
return this.checkpointService
if (!service) {
log("[Cline#initializeCheckpoints] failed to create checkpoint service, disabling checkpoints")
this.enableCheckpoints = false
return
}
service.on("initialize", ({ workspaceDir, created, duration }) => {
try {
if (created) {
log(`[Cline#initializeCheckpoints] created new shadow repo (${workspaceDir}) in ${duration}ms`)
} else {
log(
`[Cline#initializeCheckpoints] found existing shadow repo (${workspaceDir}) in ${duration}ms`,
)
}
this.checkpointService = service
this.checkpointSave()
} catch (err) {
log("[Cline#initializeCheckpoints] caught error in on('initialize'), disabling checkpoints")
this.enableCheckpoints = false
}
})
service.on("checkpoint", ({ isFirst, fromHash: from, toHash: to }) => {
try {
log(`[Cline#initializeCheckpoints] ${isFirst ? "initial" : "incremental"} checkpoint saved: ${to}`)
this.providerRef.deref()?.postMessageToWebview({ type: "currentCheckpointUpdated", text: to })
this.say("checkpoint_saved", to, undefined, undefined, { isFirst, from, to }).catch((e) =>
console.error("Error saving checkpoint message:", e),
)
} catch (err) {
log("[Cline#initializeCheckpoints] caught error in on('checkpoint'), disabling checkpoints")
this.enableCheckpoints = false
}
})
service.initShadowGit()
} catch (err) {
log("[Cline#initializeCheckpoints] caught error in initializeCheckpoints(), disabling checkpoints")
this.enableCheckpoints = false
}
}
public async checkpointDiff({
@ -3398,7 +3445,7 @@ export class Cline {
commitHash: string
mode: "full" | "checkpoint"
}) {
if (!this.enableCheckpoints) {
if (!this.checkpointService || !this.enableCheckpoints) {
return
}
@ -3414,8 +3461,7 @@ export class Cline {
}
try {
const service = await this.getCheckpointService()
const changes = await service.getDiff({ from: previousCommitHash, to: commitHash })
const changes = await this.checkpointService.getDiff({ from: previousCommitHash, to: commitHash })
if (!changes?.length) {
vscode.window.showInformationMessage("No changes found.")
@ -3441,30 +3487,13 @@ export class Cline {
}
}
public async checkpointSave({ isFirst }: { isFirst: boolean }) {
if (!this.enableCheckpoints) {
public checkpointSave() {
if (!this.checkpointService || !this.enableCheckpoints) {
return
}
try {
const service = await this.getCheckpointService()
const strategy = service.strategy
const version = service.version
const commit = await service.saveCheckpoint(`Task: ${this.taskId}, Time: ${Date.now()}`)
const fromHash = service.baseHash
const toHash = isFirst ? commit?.commit || fromHash : commit?.commit
if (toHash) {
await this.providerRef.deref()?.postMessageToWebview({ type: "currentCheckpointUpdated", text: toHash })
const checkpoint = { isFirst, from: fromHash, to: toHash, strategy, version }
await this.say("checkpoint_saved", toHash, undefined, undefined, checkpoint)
}
} catch (err) {
this.providerRef.deref()?.log("[checkpointSave] disabling checkpoints for this task")
this.enableCheckpoints = false
}
// Start the checkpoint process in the background.
this.checkpointService.saveCheckpoint(`Task: ${this.taskId}, Time: ${Date.now()}`)
}
public async checkpointRestore({
@ -3476,7 +3505,7 @@ export class Cline {
commitHash: string
mode: "preview" | "restore"
}) {
if (!this.enableCheckpoints) {
if (!this.checkpointService || !this.enableCheckpoints) {
return
}
@ -3487,8 +3516,7 @@ export class Cline {
}
try {
const service = await this.getCheckpointService()
await service.restoreCheckpoint(commitHash)
await this.checkpointService.restoreCheckpoint(commitHash)
await this.providerRef.deref()?.postMessageToWebview({ type: "currentCheckpointUpdated", text: commitHash })

29
src/services/checkpoints/CheckpointServiceFactory.ts
View file

@ -1,29 +0,0 @@
import { LocalCheckpointService, LocalCheckpointServiceOptions } from "./LocalCheckpointService"
import { ShadowCheckpointService, ShadowCheckpointServiceOptions } from "./ShadowCheckpointService"
export type CreateCheckpointServiceFactoryOptions =
| {
strategy: "local"
options: LocalCheckpointServiceOptions
}
| {
strategy: "shadow"
options: ShadowCheckpointServiceOptions
}
type CheckpointServiceType<T extends CreateCheckpointServiceFactoryOptions> = T extends { strategy: "local" }
? LocalCheckpointService
: T extends { strategy: "shadow" }
? ShadowCheckpointService
: never
export class CheckpointServiceFactory {
public static create<T extends CreateCheckpointServiceFactoryOptions>(options: T): CheckpointServiceType<T> {
switch (options.strategy) {
case "local":
return LocalCheckpointService.create(options.options) as any
case "shadow":
return ShadowCheckpointService.create(options.options) as any
}
}
}

440
src/services/checkpoints/LocalCheckpointService.ts
View file

@ -1,440 +0,0 @@
import fs from "fs/promises"
import { existsSync } from "fs"
import path from "path"
import simpleGit, { SimpleGit, CleanOptions } from "simple-git"
import { CheckpointStrategy, CheckpointService, CheckpointServiceOptions } from "./types"
export interface LocalCheckpointServiceOptions extends CheckpointServiceOptions {}
/**
* 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
* under the hood.
*
* HOW IT WORKS
*
* Two branches are used:
* - A main branch for normal operation (the branch you are currently on).
* - A hidden branch for storing checkpoints.
*
* Saving a checkpoint:
* - 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.
* - 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
* `git restore` and `git clean`.
*
* This approach allows for:
* - 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
*
* - Git must be installed.
* - If the current working directory is not a Git repository, we will
* initialize a new one with a .gitkeep file.
* - If you manually edit files and then restore a checkpoint, the changes
* will be lost. Addressing this adds some complexity to the implementation
* and it's not clear whether it's worth it.
*/
export class LocalCheckpointService implements CheckpointService {
private static readonly USER_NAME = "Roo Code"
private static readonly USER_EMAIL = "support@roocode.com"
private static readonly CHECKPOINT_BRANCH = "roo-code-checkpoints"
private static readonly STASH_BRANCH = "roo-code-stash"
public readonly strategy: CheckpointStrategy = "local"
public readonly version = 1
public get baseHash() {
return this._baseHash
}
constructor(
public readonly taskId: string,
public readonly git: SimpleGit,
public readonly workspaceDir: string,
private readonly mainBranch: string,
private _baseHash: string,
private readonly hiddenBranch: string,
private readonly log: (message: string) => void,
) {}
private async ensureBranch(expectedBranch: string) {
const branch = await this.git.revparse(["--abbrev-ref", "HEAD"])
if (branch.trim() !== expectedBranch) {
throw new Error(`Git branch mismatch: expected '${expectedBranch}' but found '${branch}'`)
}
}
public async getDiff({ from, to }: { from?: string; to?: string }) {
const result = []
if (!from) {
from = this.baseHash
}
const { files } = await this.git.diffSummary([`${from}..${to}`])
for (const file of files.filter((f) => !f.binary)) {
const relPath = file.file
const absPath = path.join(this.workspaceDir, relPath)
const before = await this.git.show([`${from}:${relPath}`]).catch(() => "")
const after = to
? await this.git.show([`${to}:${relPath}`]).catch(() => "")
: await fs.readFile(absPath, "utf8").catch(() => "")
result.push({
paths: { relative: relPath, absolute: absPath },
content: { before, after },
})
}
return result
}
private async restoreMain({
branch,
stashSha,
force = false,
}: {
branch: string
stashSha: string
force?: boolean
}) {
let currentBranch = await this.git.revparse(["--abbrev-ref", "HEAD"])
if (currentBranch !== this.mainBranch) {
if (force) {
try {
await this.git.checkout(["-f", this.mainBranch])
} catch (err) {
this.log(
`[restoreMain] failed to force checkout ${this.mainBranch}: ${err instanceof Error ? err.message : String(err)}`,
)
}
} else {
try {
await this.git.checkout(this.mainBranch)
} catch (err) {
this.log(
`[restoreMain] failed to checkout ${this.mainBranch}: ${err instanceof Error ? err.message : String(err)}`,
)
// Escalate to a forced checkout if we can't checkout the
// main branch under normal circumstances.
currentBranch = await this.git.revparse(["--abbrev-ref", "HEAD"])
if (currentBranch !== this.mainBranch) {
await this.git.checkout(["-f", this.mainBranch]).catch(() => {})
}
}
}
}
currentBranch = await this.git.revparse(["--abbrev-ref", "HEAD"])
if (currentBranch !== this.mainBranch) {
throw new Error(`Unable to restore ${this.mainBranch}`)
}
if (stashSha) {
this.log(`[restoreMain] applying stash ${stashSha}`)
try {
await this.git.raw(["stash", "apply", "--index", stashSha])
} catch (err) {
this.log(`[restoreMain] Failed to apply stash: ${err instanceof Error ? err.message : String(err)}`)
}
}
this.log(`[restoreMain] restoring from ${branch} branch`)
try {
await this.git.raw(["restore", "--source", branch, "--worktree", "--", "."])
} catch (err) {
this.log(`[restoreMain] Failed to restore branch: ${err instanceof Error ? err.message : String(err)}`)
}
}
public async saveCheckpoint(message: string) {
const startTime = Date.now()
await this.ensureBranch(this.mainBranch)
const stashSha = (await this.git.raw(["stash", "create"])).trim()
const latestSha = await this.git.revparse([this.hiddenBranch])
/**
* PHASE: Create stash
* Mutations:
* - Create branch
* - Change branch
*/
const stashBranch = `${LocalCheckpointService.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"])
} catch (err) {
this.log(
`[saveCheckpoint] failed in stage stash phase: ${err instanceof Error ? err.message : String(err)}`,
)
await this.restoreMain({ branch: stashBranch, stashSha, force: true })
await this.git.branch(["-D", stashBranch]).catch(() => {})
throw err
}
/**
* Phase: Commit stash
* Mutations:
* - Commit stash
* - Change branch
* Recovery:
* - UNDO: Create branch
* - UNDO: Change branch
*/
let stashCommit
try {
stashCommit = await this.git.commit(message, undefined, { "--no-verify": null })
this.log(`[saveCheckpoint] stashCommit: ${message} -> ${JSON.stringify(stashCommit)}`)
} catch (err) {
this.log(
`[saveCheckpoint] failed in stash commit phase: ${err instanceof Error ? err.message : String(err)}`,
)
await this.restoreMain({ branch: stashBranch, stashSha, force: true })
await this.git.branch(["-D", stashBranch]).catch(() => {})
throw err
}
if (!stashCommit) {
this.log("[saveCheckpoint] no stash commit")
await this.restoreMain({ branch: stashBranch, stashSha })
await this.git.branch(["-D", stashBranch])
return undefined
}
/**
* PHASE: Diff
* Mutations:
* - Checkout hidden branch
* Recovery:
* - UNDO: Create branch
* - UNDO: Change branch
* - UNDO: Commit stash
*/
let diff
try {
diff = await this.git.diff([latestSha, stashBranch])
} catch (err) {
this.log(`[saveCheckpoint] failed in diff phase: ${err instanceof Error ? err.message : String(err)}`)
await this.restoreMain({ branch: stashBranch, stashSha, force: true })
await this.git.branch(["-D", stashBranch]).catch(() => {})
throw err
}
if (!diff) {
this.log("[saveCheckpoint] no diff")
await this.restoreMain({ branch: stashBranch, stashSha })
await this.git.branch(["-D", stashBranch])
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) {
this.log(`[saveCheckpoint] failed in reset phase: ${err instanceof Error ? err.message : String(err)}`)
await this.restoreMain({ branch: stashBranch, stashSha, force: true })
await this.git.branch(["-D", stashBranch]).catch(() => {})
throw err
}
/**
* 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 = ""
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.workspaceDir, ".git/CHERRY_PICK_HEAD"))) {
await this.git.raw(["commit", "--allow-empty", "--no-edit"])
} else {
throw err
}
}
commit = await this.git.revparse(["HEAD"])
this.log(`[saveCheckpoint] cherry-pick commit = ${commit}`)
} catch (err) {
this.log(
`[saveCheckpoint] failed in cherry pick phase: ${err instanceof Error ? err.message : String(err)}`,
)
await this.git.reset(["--hard", latestSha]).catch(() => {})
await this.restoreMain({ branch: stashBranch, stashSha, force: true })
await this.git.branch(["-D", stashBranch]).catch(() => {})
throw err
}
await this.restoreMain({ branch: stashBranch, stashSha })
await this.git.branch(["-D", stashBranch])
// We've gotten reports that checkpoints can be slow in some cases, so
// we'll log the duration of the checkpoint save.
const duration = Date.now() - startTime
this.log(`[saveCheckpoint] saved checkpoint ${commit} in ${duration}ms`)
return { commit }
}
public async restoreCheckpoint(commitHash: string) {
const startTime = Date.now()
await this.ensureBranch(this.mainBranch)
await this.git.clean([CleanOptions.FORCE, CleanOptions.RECURSIVE])
await this.git.raw(["restore", "--source", commitHash, "--worktree", "--", "."])
const duration = Date.now() - startTime
this.log(`[restoreCheckpoint] restored checkpoint ${commitHash} in ${duration}ms`)
}
public static async create({ taskId, workspaceDir, log = console.log }: LocalCheckpointServiceOptions) {
const git = simpleGit(workspaceDir)
const version = await git.version()
if (!version?.installed) {
throw new Error(`Git is not installed. Please install Git if you wish to use checkpoints.`)
}
if (!workspaceDir || !existsSync(workspaceDir)) {
throw new Error(`Base directory is not set or does not exist.`)
}
const { currentBranch, currentSha, hiddenBranch } = await LocalCheckpointService.initRepo(git, {
taskId,
workspaceDir,
log,
})
log(
`[create] taskId = ${taskId}, workspaceDir = ${workspaceDir}, currentBranch = ${currentBranch}, currentSha = ${currentSha}, hiddenBranch = ${hiddenBranch}`,
)
return new LocalCheckpointService(taskId, git, workspaceDir, currentBranch, currentSha, hiddenBranch, log)
}
private static async initRepo(
git: SimpleGit,
{ taskId, workspaceDir, log }: Required<LocalCheckpointServiceOptions>,
) {
const isExistingRepo = existsSync(path.join(workspaceDir, ".git"))
if (!isExistingRepo) {
await git.init()
log(`[initRepo] Initialized new Git repository at ${workspaceDir}`)
}
const globalUserName = await git.getConfig("user.name", "global")
const localUserName = await git.getConfig("user.name", "local")
const userName = localUserName.value || globalUserName.value
const globalUserEmail = await git.getConfig("user.email", "global")
const localUserEmail = await git.getConfig("user.email", "local")
const userEmail = localUserEmail.value || globalUserEmail.value
// Prior versions of this service indiscriminately set the local user
// config, and it should not override the global config. To address
// this we remove the local user config if it matches the default
// user name and email and there's a global config.
if (globalUserName.value && localUserName.value === LocalCheckpointService.USER_NAME) {
await git.raw(["config", "--unset", "--local", "user.name"])
}
if (globalUserEmail.value && localUserEmail.value === LocalCheckpointService.USER_EMAIL) {
await git.raw(["config", "--unset", "--local", "user.email"])
}
// Only set user config if not already configured.
if (!userName) {
await git.addConfig("user.name", LocalCheckpointService.USER_NAME)
}
if (!userEmail) {
await git.addConfig("user.email", LocalCheckpointService.USER_EMAIL)
}
if (!isExistingRepo) {
// We need at least one file to commit, otherwise the initial
// commit will fail, unless we use the `--allow-empty` flag.
// However, using an empty commit causes problems when restoring
// the checkpoint (i.e. the `git restore` command doesn't work
// for empty commits).
await fs.writeFile(path.join(workspaceDir, ".gitkeep"), "")
await git.add(".gitkeep")
const commit = await git.commit("Initial commit")
if (!commit.commit) {
throw new Error("Failed to create initial commit")
}
log(`[initRepo] Initial commit: ${commit.commit}`)
}
const currentBranch = await git.revparse(["--abbrev-ref", "HEAD"])
const currentSha = await git.revparse(["HEAD"])
const hiddenBranch = `${LocalCheckpointService.CHECKPOINT_BRANCH}-${taskId}`
const branchSummary = await git.branch()
if (!branchSummary.all.includes(hiddenBranch)) {
await git.checkoutBranch(hiddenBranch, currentBranch)
await git.checkout(currentBranch)
}
return { currentBranch, currentSha, hiddenBranch }
}
}

135
src/services/checkpoints/ShadowCheckpointService.ts
View file

@ -5,17 +5,18 @@ import { globby } from "globby"
import simpleGit, { SimpleGit } from "simple-git"
import { GIT_DISABLED_SUFFIX, GIT_EXCLUDES } from "./constants"
import { CheckpointStrategy, CheckpointService, CheckpointServiceOptions } from "./types"
import { CheckpointService, CheckpointServiceOptions, CheckpointEventEmitter } from "./types"
export interface ShadowCheckpointServiceOptions extends CheckpointServiceOptions {
shadowDir: string
}
export class ShadowCheckpointService implements CheckpointService {
public readonly strategy: CheckpointStrategy = "shadow"
export class ShadowCheckpointService extends CheckpointEventEmitter implements CheckpointService {
public readonly version = 1
private _checkpoints: string[] = []
private _baseHash?: string
private _isInitialized = false
public get baseHash() {
return this._baseHash
@ -25,6 +26,18 @@ export class ShadowCheckpointService implements CheckpointService {
this._baseHash = value
}
public get isInitialized() {
return this._isInitialized
}
private set isInitialized(value: boolean) {
this._isInitialized = value
}
public get isDirty() {
return this._checkpoints.length > 0
}
private readonly shadowGitDir: string
private shadowGitConfigWorktree?: string
@ -35,18 +48,26 @@ export class ShadowCheckpointService implements CheckpointService {
public readonly workspaceDir: string,
private readonly log: (message: string) => void,
) {
super()
this.shadowGitDir = path.join(this.shadowDir, "tasks", this.taskId, "checkpoints", ".git")
}
private async initShadowGit() {
public async initShadowGit() {
if (this.isInitialized) {
return
}
const fileExistsAtPath = (path: string) =>
fs
.access(path)
.then(() => true)
.catch(() => false)
let created = false
const startTime = Date.now()
if (await fileExistsAtPath(this.shadowGitDir)) {
this.log(`[initShadowGit] shadow git repo already exists at ${this.shadowGitDir}`)
this.log(`[CheckpointService#initShadowGit] shadow git repo already exists at ${this.shadowGitDir}`)
const worktree = await this.getShadowGitConfigWorktree()
if (worktree !== this.workspaceDir) {
@ -57,7 +78,7 @@ export class ShadowCheckpointService implements CheckpointService {
this.baseHash = await this.git.revparse(["--abbrev-ref", "HEAD"])
} else {
this.log(`[initShadowGit] creating shadow git repo at ${this.workspaceDir}`)
this.log(`[CheckpointService#initShadowGit] creating shadow git repo at ${this.workspaceDir}`)
await this.git.init()
await this.git.addConfig("core.worktree", this.workspaceDir) // Sets the working tree to the current workspace.
@ -78,7 +99,7 @@ export class ShadowCheckpointService implements CheckpointService {
}
} catch (error) {
this.log(
`[initShadowGit] failed to read .gitattributes: ${error instanceof Error ? error.message : String(error)}`,
`[CheckpointService#initShadowGit] failed to read .gitattributes: ${error instanceof Error ? error.message : String(error)}`,
)
}
@ -93,8 +114,23 @@ export class ShadowCheckpointService implements CheckpointService {
await this.stageAll()
const { commit } = await this.git.commit("initial commit", { "--allow-empty": null })
this.baseHash = commit
this.log(`[initShadowGit] base commit is ${commit}`)
this.log(`[CheckpointService#initShadowGit] base commit is ${commit}`)
created = true
}
const duration = Date.now() - startTime
this.log(`[CheckpointService#initShadowGit] initialized shadow git in ${duration}ms`)
this.isInitialized = true
this.emit("initialize", {
type: "initialize",
workspaceDir: this.workspaceDir,
baseHash: this.baseHash,
created,
duration,
})
}
private async stageAll() {
@ -103,7 +139,9 @@ export class ShadowCheckpointService implements CheckpointService {
try {
await this.git.add(".")
} catch (error) {
this.log(`[stageAll] failed to add files to git: ${error instanceof Error ? error.message : String(error)}`)
this.log(
`[CheckpointService#stageAll] failed to add files to git: ${error instanceof Error ? error.message : String(error)}`,
)
} finally {
await this.renameNestedGitRepos(false)
}
@ -137,10 +175,12 @@ export class ShadowCheckpointService implements CheckpointService {
try {
await fs.rename(fullPath, newPath)
this.log(`${disable ? "disabled" : "enabled"} nested git repo ${gitPath}`)
this.log(
`[CheckpointService#renameNestedGitRepos] ${disable ? "disabled" : "enabled"} nested git repo ${gitPath}`,
)
} catch (error) {
this.log(
`failed to ${disable ? "disable" : "enable"} nested git repo ${gitPath}: ${error instanceof Error ? error.message : String(error)}`,
`[CheckpointService#renameNestedGitRepos] failed to ${disable ? "disable" : "enable"} nested git repo ${gitPath}: ${error instanceof Error ? error.message : String(error)}`,
)
}
}
@ -152,7 +192,7 @@ export class ShadowCheckpointService implements CheckpointService {
this.shadowGitConfigWorktree = (await this.git.getConfig("core.worktree")).value || undefined
} catch (error) {
this.log(
`[getShadowGitConfigWorktree] failed to get core.worktree: ${error instanceof Error ? error.message : String(error)}`,
`[CheckpointService#getShadowGitConfigWorktree] failed to get core.worktree: ${error instanceof Error ? error.message : String(error)}`,
)
}
}
@ -162,35 +202,63 @@ export class ShadowCheckpointService implements CheckpointService {
public async saveCheckpoint(message: string) {
try {
this.log("[CheckpointService#saveCheckpoint] starting checkpoint save")
if (!this.isInitialized) {
throw new Error("Shadow git repo not initialized")
}
const startTime = Date.now()
await this.stageAll()
const result = await this.git.commit(message)
if (result.commit) {
const duration = Date.now() - startTime
this.log(`[saveCheckpoint] saved checkpoint ${result.commit} in ${duration}ms`)
return result
} else {
return undefined
}
} catch (error) {
this.log(
`[saveCheckpoint] failed to create checkpoint: ${error instanceof Error ? error.message : String(error)}`,
)
const isFirst = this._checkpoints.length === 0
const fromHash = this._checkpoints[this._checkpoints.length - 1] ?? this.baseHash!
const toHash = result.commit ?? fromHash
this._checkpoints.push(toHash)
const duration = Date.now() - startTime
this.emit("checkpoint", { type: "checkpoint", isFirst, fromHash, toHash, duration })
return result.commit ? result : undefined
} catch (e) {
const error = e instanceof Error ? e : new Error(String(e))
this.log(`[CheckpointService#saveCheckpoint] failed to create checkpoint: ${error.message}`)
this.emit("error", { type: "error", error })
throw error
}
}
public async restoreCheckpoint(commitHash: string) {
const start = Date.now()
await this.git.clean("f", ["-d", "-f"])
await this.git.reset(["--hard", commitHash])
const duration = Date.now() - start
this.log(`[restoreCheckpoint] restored checkpoint ${commitHash} in ${duration}ms`)
try {
if (!this.isInitialized) {
throw new Error("Shadow git repo not initialized")
}
const start = Date.now()
await this.git.clean("f", ["-d", "-f"])
await this.git.reset(["--hard", commitHash])
// Remove all checkpoints after the specified commitHash.
const checkpointIndex = this._checkpoints.indexOf(commitHash)
if (checkpointIndex !== -1) {
this._checkpoints = this._checkpoints.slice(0, checkpointIndex + 1)
}
const duration = Date.now() - start
this.emit("restore", { type: "restore", commitHash, duration })
this.log(`[CheckpointService#restoreCheckpoint] restored checkpoint ${commitHash} in ${duration}ms`)
} catch (e) {
const error = e instanceof Error ? e : new Error(String(e))
this.log(`[CheckpointService#restoreCheckpoint] failed to restore checkpoint: ${error.message}`)
this.emit("error", { type: "error", error })
throw error
}
}
public async getDiff({ from, to }: { from?: string; to?: string }) {
if (!this.isInitialized) {
throw new Error("Shadow git repo not initialized")
}
const result = []
if (!from) {
@ -223,6 +291,7 @@ export class ShadowCheckpointService implements CheckpointService {
try {
await simpleGit().version()
} catch (error) {
log("[CheckpointService#create] git is not installed")
throw new Error("Git must be installed to use checkpoints.")
}
@ -241,9 +310,7 @@ export class ShadowCheckpointService implements CheckpointService {
const gitDir = path.join(checkpointsDir, ".git")
const git = simpleGit(path.dirname(gitDir))
log(`[create] taskId = ${taskId}, workspaceDir = ${workspaceDir}, shadowDir = ${shadowDir}`)
const service = new ShadowCheckpointService(taskId, git, shadowDir, workspaceDir, log)
await service.initShadowGit()
return service
log(`[CheckpointService#create] taskId = ${taskId}, workspaceDir = ${workspaceDir}, shadowDir = ${shadowDir}`)
return new ShadowCheckpointService(taskId, git, shadowDir, workspaceDir, log)
}
}

385
src/services/checkpoints/__tests__/LocalCheckpointService.test.ts
View file

@ -1,385 +0,0 @@
// npx jest src/services/checkpoints/__tests__/LocalCheckpointService.test.ts
import fs from "fs/promises"
import path from "path"
import os from "os"
import { simpleGit, SimpleGit } from "simple-git"
import { CheckpointServiceFactory } from "../CheckpointServiceFactory"
import { LocalCheckpointService } from "../LocalCheckpointService"
const tmpDir = path.join(os.tmpdir(), "test-LocalCheckpointService")
describe("LocalCheckpointService", () => {
const taskId = "test-task"
let testFile: string
let service: LocalCheckpointService
const initRepo = async ({
workspaceDir,
userName = "Roo Code",
userEmail = "support@roocode.com",
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 { testFile }
}
beforeEach(async () => {
const workspaceDir = path.join(tmpDir, `checkpoint-service-test-${Date.now()}`)
const repo = await initRepo({ workspaceDir })
testFile = repo.testFile
service = await CheckpointServiceFactory.create({
strategy: "local",
options: { taskId, workspaceDir, log: () => {} },
})
})
afterEach(async () => {
jest.restoreAllMocks()
})
afterAll(async () => {
await fs.rm(tmpDir, { recursive: true, force: true })
})
describe("getDiff", () => {
it("returns the correct diff between commits", async () => {
await fs.writeFile(testFile, "Ahoy, world!")
const commit1 = await service.saveCheckpoint("First checkpoint")
expect(commit1?.commit).toBeTruthy()
await fs.writeFile(testFile, "Goodbye, world!")
const commit2 = await service.saveCheckpoint("Second checkpoint")
expect(commit2?.commit).toBeTruthy()
const diff1 = await service.getDiff({ to: commit1!.commit })
expect(diff1).toHaveLength(1)
expect(diff1[0].paths.relative).toBe("test.txt")
expect(diff1[0].paths.absolute).toBe(testFile)
expect(diff1[0].content.before).toBe("Hello, world!")
expect(diff1[0].content.after).toBe("Ahoy, world!")
const diff2 = await service.getDiff({ to: commit2!.commit })
expect(diff2).toHaveLength(1)
expect(diff2[0].paths.relative).toBe("test.txt")
expect(diff2[0].paths.absolute).toBe(testFile)
expect(diff2[0].content.before).toBe("Hello, world!")
expect(diff2[0].content.after).toBe("Goodbye, world!")
const diff12 = await service.getDiff({ from: commit1!.commit, to: commit2!.commit })
expect(diff12).toHaveLength(1)
expect(diff12[0].paths.relative).toBe("test.txt")
expect(diff12[0].paths.absolute).toBe(testFile)
expect(diff12[0].content.before).toBe("Ahoy, world!")
expect(diff12[0].content.after).toBe("Goodbye, world!")
})
it("handles new files in diff", async () => {
const newFile = path.join(service.workspaceDir, "new.txt")
await fs.writeFile(newFile, "New file content")
const commit = await service.saveCheckpoint("Add new file")
expect(commit?.commit).toBeTruthy()
const changes = await service.getDiff({ to: commit!.commit })
const change = changes.find((c) => c.paths.relative === "new.txt")
expect(change).toBeDefined()
expect(change?.content.before).toBe("")
expect(change?.content.after).toBe("New file content")
})
it("handles deleted files in diff", async () => {
const fileToDelete = path.join(service.workspaceDir, "new.txt")
await fs.writeFile(fileToDelete, "New file content")
const commit1 = await service.saveCheckpoint("Add file")
expect(commit1?.commit).toBeTruthy()
await fs.unlink(fileToDelete)
const commit2 = await service.saveCheckpoint("Delete file")
expect(commit2?.commit).toBeTruthy()
const changes = await service.getDiff({ from: commit1!.commit, to: commit2!.commit })
const change = changes.find((c) => c.paths.relative === "new.txt")
expect(change).toBeDefined()
expect(change!.content.before).toBe("New file content")
expect(change!.content.after).toBe("")
})
})
describe("saveCheckpoint", () => {
it("creates a checkpoint if there are pending changes", async () => {
await fs.writeFile(testFile, "Ahoy, world!")
const commit1 = await service.saveCheckpoint("First checkpoint")
expect(commit1?.commit).toBeTruthy()
const details1 = await service.git.show([commit1!.commit])
expect(details1).toContain("-Hello, world!")
expect(details1).toContain("+Ahoy, world!")
await fs.writeFile(testFile, "Hola, world!")
const commit2 = await service.saveCheckpoint("Second checkpoint")
expect(commit2?.commit).toBeTruthy()
const details2 = await service.git.show([commit2!.commit])
expect(details2).toContain("-Hello, world!")
expect(details2).toContain("+Hola, world!")
// Switch to checkpoint 1.
await service.restoreCheckpoint(commit1!.commit)
expect(await fs.readFile(testFile, "utf-8")).toBe("Ahoy, world!")
// Switch to checkpoint 2.
await service.restoreCheckpoint(commit2!.commit)
expect(await fs.readFile(testFile, "utf-8")).toBe("Hola, world!")
// Switch back to initial commit.
expect(service.baseHash).toBeTruthy()
await service.restoreCheckpoint(service.baseHash!)
expect(await fs.readFile(testFile, "utf-8")).toBe("Hello, world!")
})
it("preserves workspace and index state after saving checkpoint", async () => {
// Create three files with different states: staged, unstaged, and mixed.
const unstagedFile = path.join(service.workspaceDir, "unstaged.txt")
const stagedFile = path.join(service.workspaceDir, "staged.txt")
const mixedFile = path.join(service.workspaceDir, "mixed.txt")
await fs.writeFile(unstagedFile, "Initial unstaged")
await fs.writeFile(stagedFile, "Initial staged")
await fs.writeFile(mixedFile, "Initial mixed")
await service.git.add(["."])
const result = await service.git.commit("Add initial files")
expect(result?.commit).toBeTruthy()
await fs.writeFile(unstagedFile, "Modified unstaged")
await fs.writeFile(stagedFile, "Modified staged")
await service.git.add([stagedFile])
await fs.writeFile(mixedFile, "Modified mixed - staged")
await service.git.add([mixedFile])
await fs.writeFile(mixedFile, "Modified mixed - unstaged")
// Save checkpoint.
const commit = await service.saveCheckpoint("Test checkpoint")
expect(commit?.commit).toBeTruthy()
// Verify workspace state is preserved.
const status = await service.git.status()
// All files should be modified.
expect(status.modified).toContain("unstaged.txt")
expect(status.modified).toContain("staged.txt")
expect(status.modified).toContain("mixed.txt")
// Only staged and mixed files should be staged.
expect(status.staged).not.toContain("unstaged.txt")
expect(status.staged).toContain("staged.txt")
expect(status.staged).toContain("mixed.txt")
// Verify file contents.
expect(await fs.readFile(unstagedFile, "utf-8")).toBe("Modified unstaged")
expect(await fs.readFile(stagedFile, "utf-8")).toBe("Modified staged")
expect(await fs.readFile(mixedFile, "utf-8")).toBe("Modified mixed - unstaged")
// Verify staged changes (--cached shows only staged changes).
const stagedDiff = await service.git.diff(["--cached", "mixed.txt"])
expect(stagedDiff).toContain("-Initial mixed")
expect(stagedDiff).toContain("+Modified mixed - staged")
// Verify unstaged changes (shows working directory changes).
const unstagedDiff = await service.git.diff(["mixed.txt"])
expect(unstagedDiff).toContain("-Modified mixed - staged")
expect(unstagedDiff).toContain("+Modified mixed - unstaged")
})
it("does not create a checkpoint if there are no pending changes", async () => {
const commit0 = await service.saveCheckpoint("Zeroth checkpoint")
expect(commit0?.commit).toBeFalsy()
await fs.writeFile(testFile, "Ahoy, world!")
const commit1 = await service.saveCheckpoint("First checkpoint")
expect(commit1?.commit).toBeTruthy()
const commit2 = await service.saveCheckpoint("Second checkpoint")
expect(commit2?.commit).toBeFalsy()
})
it("includes untracked files in checkpoints", async () => {
// Create an untracked file.
const untrackedFile = path.join(service.workspaceDir, "untracked.txt")
await fs.writeFile(untrackedFile, "I am untracked!")
// Save a checkpoint with the untracked file.
const commit1 = await service.saveCheckpoint("Checkpoint with untracked file")
expect(commit1?.commit).toBeTruthy()
// Verify the untracked file was included in the checkpoint.
const details = await service.git.show([commit1!.commit])
expect(details).toContain("+I am untracked!")
// Create another checkpoint with a different state.
await fs.writeFile(testFile, "Changed tracked file")
const commit2 = await service.saveCheckpoint("Second checkpoint")
expect(commit2?.commit).toBeTruthy()
// Restore first checkpoint and verify untracked file is preserved.
await service.restoreCheckpoint(commit1!.commit)
expect(await fs.readFile(untrackedFile, "utf-8")).toBe("I am untracked!")
expect(await fs.readFile(testFile, "utf-8")).toBe("Hello, world!")
// Restore second checkpoint and verify untracked file remains (since
// restore preserves untracked files)
await service.restoreCheckpoint(commit2!.commit)
expect(await fs.readFile(untrackedFile, "utf-8")).toBe("I am untracked!")
expect(await fs.readFile(testFile, "utf-8")).toBe("Changed tracked file")
})
it("throws if we're on the wrong branch", async () => {
// Create and switch to a feature branch.
const currentBranch = await service.git.revparse(["--abbrev-ref", "HEAD"])
await service.git.checkoutBranch("feature", currentBranch)
// Attempt to save checkpoint from feature branch.
await expect(service.saveCheckpoint("test")).rejects.toThrow(
`Git branch mismatch: expected '${currentBranch}' but found 'feature'`,
)
// Attempt to restore checkpoint from feature branch.
expect(service.baseHash).toBeTruthy()
await expect(service.restoreCheckpoint(service.baseHash!)).rejects.toThrow(
`Git branch mismatch: expected '${currentBranch}' but found 'feature'`,
)
})
it("cleans up staged files if a commit fails", async () => {
await fs.writeFile(testFile, "Changed content")
// Mock git commit to simulate failure.
jest.spyOn(service.git, "commit").mockRejectedValue(new Error("Simulated commit failure"))
// Attempt to save checkpoint.
await expect(service.saveCheckpoint("test")).rejects.toThrow("Simulated commit failure")
// Verify files are unstaged.
const status = await service.git.status()
expect(status.staged).toHaveLength(0)
})
it("handles file deletions correctly", async () => {
await fs.writeFile(testFile, "I am tracked!")
const untrackedFile = path.join(service.workspaceDir, "new.txt")
await fs.writeFile(untrackedFile, "I am untracked!")
const commit1 = await service.saveCheckpoint("First checkpoint")
expect(commit1?.commit).toBeTruthy()
await fs.unlink(testFile)
await fs.unlink(untrackedFile)
const commit2 = await service.saveCheckpoint("Second checkpoint")
expect(commit2?.commit).toBeTruthy()
// Verify files are gone.
await expect(fs.readFile(testFile, "utf-8")).rejects.toThrow()
await expect(fs.readFile(untrackedFile, "utf-8")).rejects.toThrow()
// Restore first checkpoint.
await service.restoreCheckpoint(commit1!.commit)
expect(await fs.readFile(testFile, "utf-8")).toBe("I am tracked!")
expect(await fs.readFile(untrackedFile, "utf-8")).toBe("I am untracked!")
// Restore second checkpoint.
await service.restoreCheckpoint(commit2!.commit)
await expect(fs.readFile(testFile, "utf-8")).rejects.toThrow()
await expect(fs.readFile(untrackedFile, "utf-8")).rejects.toThrow()
})
})
describe("create", () => {
it("initializes a git repository if one does not already exist", async () => {
const workspaceDir = path.join(tmpDir, `checkpoint-service-test2-${Date.now()}`)
await fs.mkdir(workspaceDir)
const newTestFile = path.join(workspaceDir, "test.txt")
await fs.writeFile(newTestFile, "Hello, world!")
// Ensure the git repository was initialized.
const gitDir = path.join(workspaceDir, ".git")
await expect(fs.stat(gitDir)).rejects.toThrow()
const newService = await LocalCheckpointService.create({ taskId, workspaceDir, log: () => {} })
expect(await fs.stat(gitDir)).toBeTruthy()
// Save a checkpoint: Hello, world!
const commit1 = await newService.saveCheckpoint("Hello, world!")
expect(commit1?.commit).toBeTruthy()
expect(await fs.readFile(newTestFile, "utf-8")).toBe("Hello, world!")
// Restore initial commit; the file should no longer exist.
expect(newService.baseHash).toBeTruthy()
await newService.restoreCheckpoint(newService.baseHash!)
await expect(fs.access(newTestFile)).rejects.toThrow()
// Restore to checkpoint 1; the file should now exist.
await newService.restoreCheckpoint(commit1!.commit)
expect(await fs.readFile(newTestFile, "utf-8")).toBe("Hello, world!")
// Save a new checkpoint: Ahoy, world!
await fs.writeFile(newTestFile, "Ahoy, world!")
const commit2 = await newService.saveCheckpoint("Ahoy, world!")
expect(commit2?.commit).toBeTruthy()
expect(await fs.readFile(newTestFile, "utf-8")).toBe("Ahoy, world!")
// Restore "Hello, world!"
await newService.restoreCheckpoint(commit1!.commit)
expect(await fs.readFile(newTestFile, "utf-8")).toBe("Hello, world!")
// Restore "Ahoy, world!"
await newService.restoreCheckpoint(commit2!.commit)
expect(await fs.readFile(newTestFile, "utf-8")).toBe("Ahoy, world!")
// Restore initial commit.
expect(newService.baseHash).toBeTruthy()
await newService.restoreCheckpoint(newService.baseHash!)
await expect(fs.access(newTestFile)).rejects.toThrow()
await fs.rm(newService.workspaceDir, { recursive: true, force: true })
})
it("respects existing git user configuration", async () => {
const workspaceDir = path.join(tmpDir, `checkpoint-service-test-config2-${Date.now()}`)
const userName = "Custom User"
const userEmail = "custom@example.com"
await initRepo({ workspaceDir, userName, userEmail })
const newService = await LocalCheckpointService.create({ taskId, workspaceDir, log: () => {} })
expect((await newService.git.getConfig("user.name")).value).toBe(userName)
expect((await newService.git.getConfig("user.email")).value).toBe(userEmail)
await fs.rm(workspaceDir, { recursive: true, force: true })
})
})
})

176
src/services/checkpoints/__tests__/ShadowCheckpointService.test.ts
View file

@ -3,11 +3,11 @@
import fs from "fs/promises"
import path from "path"
import os from "os"
import { EventEmitter } from "events"
import { simpleGit, SimpleGit } from "simple-git"
import { ShadowCheckpointService } from "../ShadowCheckpointService"
import { CheckpointServiceFactory } from "../CheckpointServiceFactory"
jest.mock("globby", () => ({
globby: jest.fn().mockResolvedValue([]),
@ -63,13 +63,10 @@ describe("ShadowCheckpointService", () => {
const repo = await initRepo({ workspaceDir })
testFile = repo.testFile
service = await CheckpointServiceFactory.create({
strategy: "shadow",
options: { taskId, shadowDir, workspaceDir, log: () => {} },
})
workspaceGit = repo.git
service = await ShadowCheckpointService.create({ taskId, shadowDir, workspaceDir, log: () => {} })
await service.initShadowGit()
})
afterEach(async () => {
@ -313,6 +310,7 @@ describe("ShadowCheckpointService", () => {
const gitDir = path.join(shadowDir, "tasks", taskId, "checkpoints", ".git")
await expect(fs.stat(gitDir)).rejects.toThrow()
const newService = await ShadowCheckpointService.create({ taskId, shadowDir, workspaceDir, log: () => {} })
await newService.initShadowGit()
expect(await fs.stat(gitDir)).toBeTruthy()
// Save a new checkpoint: Ahoy, world!
@ -333,4 +331,168 @@ describe("ShadowCheckpointService", () => {
await fs.rm(newService.workspaceDir, { recursive: true, force: true })
})
})
describe("events", () => {
it("emits initialize event when service is created", async () => {
const shadowDir = path.join(tmpDir, `shadow-event-test-${Date.now()}`)
const workspaceDir = path.join(tmpDir, `workspace-event-test-${Date.now()}`)
await fs.mkdir(workspaceDir, { recursive: true })
const newTestFile = path.join(workspaceDir, "test.txt")
await fs.writeFile(newTestFile, "Testing events!")
// Create a mock implementation of emit to track events.
const emitSpy = jest.spyOn(EventEmitter.prototype, "emit")
// Create the service - this will trigger the initialize event.
const newService = await ShadowCheckpointService.create({ taskId, shadowDir, workspaceDir, log: () => {} })
await newService.initShadowGit()
// Find the initialize event in the emit calls.
let initializeEvent = null
for (let i = 0; i < emitSpy.mock.calls.length; i++) {
const call = emitSpy.mock.calls[i]
if (call[0] === "initialize") {
initializeEvent = call[1]
break
}
}
// Restore the spy.
emitSpy.mockRestore()
// Verify the event was emitted with the correct data.
expect(initializeEvent).not.toBeNull()
expect(initializeEvent.type).toBe("initialize")
expect(initializeEvent.workspaceDir).toBe(workspaceDir)
expect(initializeEvent.baseHash).toBeTruthy()
expect(typeof initializeEvent.created).toBe("boolean")
expect(typeof initializeEvent.duration).toBe("number")
// Verify the event was emitted with the correct data.
expect(initializeEvent).not.toBeNull()
expect(initializeEvent.type).toBe("initialize")
expect(initializeEvent.workspaceDir).toBe(workspaceDir)
expect(initializeEvent.baseHash).toBeTruthy()
expect(typeof initializeEvent.created).toBe("boolean")
expect(typeof initializeEvent.duration).toBe("number")
// Clean up.
await fs.rm(shadowDir, { recursive: true, force: true })
await fs.rm(workspaceDir, { recursive: true, force: true })
})
it("emits checkpoint event when saving checkpoint", async () => {
const checkpointHandler = jest.fn()
service.on("checkpoint", checkpointHandler)
await fs.writeFile(testFile, "Changed content for checkpoint event test")
const result = await service.saveCheckpoint("Test checkpoint event")
expect(result?.commit).toBeDefined()
expect(checkpointHandler).toHaveBeenCalledTimes(1)
const eventData = checkpointHandler.mock.calls[0][0]
expect(eventData.type).toBe("checkpoint")
expect(eventData.toHash).toBeDefined()
expect(eventData.toHash).toBe(result!.commit)
expect(typeof eventData.duration).toBe("number")
})
it("emits restore event when restoring checkpoint", async () => {
// First create a checkpoint to restore.
await fs.writeFile(testFile, "Content for restore test")
const commit = await service.saveCheckpoint("Checkpoint for restore test")
expect(commit?.commit).toBeTruthy()
// Change the file again.
await fs.writeFile(testFile, "Changed after checkpoint")
// Setup restore event listener.
const restoreHandler = jest.fn()
service.on("restore", restoreHandler)
// Restore the checkpoint.
await service.restoreCheckpoint(commit!.commit)
// Verify the event was emitted.
expect(restoreHandler).toHaveBeenCalledTimes(1)
const eventData = restoreHandler.mock.calls[0][0]
expect(eventData.type).toBe("restore")
expect(eventData.commitHash).toBe(commit!.commit)
expect(typeof eventData.duration).toBe("number")
// Verify the file was actually restored.
expect(await fs.readFile(testFile, "utf-8")).toBe("Content for restore test")
})
it("emits error event when an error occurs", async () => {
const errorHandler = jest.fn()
service.on("error", errorHandler)
// Force an error by providing an invalid commit hash.
const invalidCommitHash = "invalid-commit-hash"
// Try to restore an invalid checkpoint.
try {
await service.restoreCheckpoint(invalidCommitHash)
} catch (error) {
// Expected to throw, we're testing the event emission.
}
// Verify the error event was emitted.
expect(errorHandler).toHaveBeenCalledTimes(1)
const eventData = errorHandler.mock.calls[0][0]
expect(eventData.type).toBe("error")
expect(eventData.error).toBeInstanceOf(Error)
})
it("supports multiple event listeners for the same event", async () => {
const checkpointHandler1 = jest.fn()
const checkpointHandler2 = jest.fn()
service.on("checkpoint", checkpointHandler1)
service.on("checkpoint", checkpointHandler2)
await fs.writeFile(testFile, "Content for multiple listeners test")
const result = await service.saveCheckpoint("Testing multiple listeners")
// Verify both handlers were called with the same event data.
expect(checkpointHandler1).toHaveBeenCalledTimes(1)
expect(checkpointHandler2).toHaveBeenCalledTimes(1)
const eventData1 = checkpointHandler1.mock.calls[0][0]
const eventData2 = checkpointHandler2.mock.calls[0][0]
expect(eventData1).toEqual(eventData2)
expect(eventData1.type).toBe("checkpoint")
expect(eventData1.toHash).toBe(result?.commit)
})
it("allows removing event listeners", async () => {
const checkpointHandler = jest.fn()
// Add the listener.
service.on("checkpoint", checkpointHandler)
// Make a change and save a checkpoint.
await fs.writeFile(testFile, "Content for remove listener test - part 1")
await service.saveCheckpoint("Testing listener - part 1")
// Verify handler was called.
expect(checkpointHandler).toHaveBeenCalledTimes(1)
checkpointHandler.mockClear()
// Remove the listener.
service.off("checkpoint", checkpointHandler)
// Make another change and save a checkpoint.
await fs.writeFile(testFile, "Content for remove listener test - part 2")
await service.saveCheckpoint("Testing listener - part 2")
// Verify handler was not called after being removed.
expect(checkpointHandler).not.toHaveBeenCalled()
})
})
})

2
src/services/checkpoints/index.ts
View file

@ -1,2 +0,0 @@
export * from "./types"
export * from "./CheckpointServiceFactory"

39
src/services/checkpoints/types.ts
View file

@ -1,3 +1,4 @@
import EventEmitter from "events"
import { CommitResult } from "simple-git"
export type CheckpointResult = Partial<CommitResult> & Pick<CommitResult, "commit">
@ -13,15 +14,12 @@ export type CheckpointDiff = {
}
}
export type CheckpointStrategy = "local" | "shadow"
export interface CheckpointService {
saveCheckpoint(message: string): Promise<CheckpointResult | undefined>
restoreCheckpoint(commit: string): Promise<void>
getDiff(range: { from?: string; to?: string }): Promise<CheckpointDiff[]>
workspaceDir: string
baseHash?: string
strategy: CheckpointStrategy
version: number
}
@ -30,3 +28,38 @@ export interface CheckpointServiceOptions {
workspaceDir: string
log?: (message: string) => void
}
/**
* EventEmitter
*/
export interface CheckpointEventMap {
initialize: { type: "initialize"; workspaceDir: string; baseHash: string; created: boolean; duration: number }
checkpoint: {
type: "checkpoint"
isFirst: boolean
fromHash: string
toHash: string
duration: number
}
restore: { type: "restore"; commitHash: string; duration: number }
error: { type: "error"; error: Error }
}
export class CheckpointEventEmitter extends EventEmitter {
override emit<K extends keyof CheckpointEventMap>(event: K, data: CheckpointEventMap[K]): boolean {
return super.emit(event, data)
}
override on<K extends keyof CheckpointEventMap>(event: K, listener: (data: CheckpointEventMap[K]) => void): this {
return super.on(event, listener)
}
override off<K extends keyof CheckpointEventMap>(event: K, listener: (data: CheckpointEventMap[K]) => void): this {
return super.off(event, listener)
}
override once<K extends keyof CheckpointEventMap>(event: K, listener: (data: CheckpointEventMap[K]) => void): this {
return super.once(event, listener)
}
}

4
webview-ui/src/components/chat/checkpoints/CheckpointSaved.tsx
View file

@ -3,8 +3,6 @@ import { useMemo } from "react"
import { CheckpointMenu } from "./CheckpointMenu"
import { checkpointSchema } from "./schema"
const REQUIRED_VERSION = 1
type CheckpointSavedProps = {
ts: number
commitHash: string
@ -22,7 +20,7 @@ export const CheckpointSaved = ({ checkpoint, ...props }: CheckpointSavedProps)
const result = checkpointSchema.safeParse(checkpoint)
if (!result.success || result.data.version < REQUIRED_VERSION) {
if (!result.success) {
return undefined
}

2
webview-ui/src/components/chat/checkpoints/schema.ts
View file

@ -4,8 +4,6 @@ export const checkpointSchema = z.object({
isFirst: z.boolean(),
from: z.string(),
to: z.string(),
strategy: z.enum(["local", "shadow"]),
version: z.number(),
})
export type Checkpoint = z.infer<typeof checkpointSchema>