fix: prevent data loss in checkpoint restoration (fixes #6209)

src/services/checkpoints/ShadowCheckpointService.ts

@@ -64,6 +64,18 @@ export abstract class ShadowCheckpointService extends EventEmitter {
 			throw new Error("Shadow git repo already initialized")
 		}
 
+		// Additional validation to ensure workspace directory is valid
+		if (!this.workspaceDir || this.workspaceDir.trim() === "") {
+			throw new Error("Invalid workspace directory: empty or undefined")
+		}
+
+		// Ensure workspace directory exists
+		try {
+			await fs.access(this.workspaceDir)
+		} catch (error) {
+			throw new Error(`Workspace directory does not exist: ${this.workspaceDir}`)
+		}
+
 		const hasNestedGitRepos = await this.hasNestedGitRepositories()
 
 		if (hasNestedGitRepos) {
@@ -85,12 +97,26 @@ export abstract class ShadowCheckpointService extends EventEmitter {
 			this.log(`[${this.constructor.name}#initShadowGit] shadow git repo already exists at ${this.dotGitDir}`)
 			const worktree = await this.getShadowGitConfigWorktree(git)
 
-			if (worktree !== this.workspaceDir) {
+			if (!worktree) {
+				// If worktree is not set, try to repair it
+				this.log(
+					`[${this.constructor.name}#initShadowGit] worktree not configured, setting to ${this.workspaceDir}`,
+				)
+				await git.addConfig("core.worktree", this.workspaceDir)
+				this.shadowGitConfigWorktree = this.workspaceDir
+			} else if (worktree !== this.workspaceDir) {
 				throw new Error(
 					`Checkpoints can only be used in the original workspace: ${worktree} !== ${this.workspaceDir}`,
 				)
 			}
 
+			// Validate that the shadow git repo is not corrupted
+			try {
+				await git.status()
+			} catch (statusError) {
+				throw new Error(`Shadow git repository appears to be corrupted: ${statusError.message}`)
+			}
+
 			await this.writeExcludeFile()
 			this.baseHash = await git.revparse(["HEAD"])
 		} else {
@@ -246,10 +272,30 @@ export abstract class ShadowCheckpointService extends EventEmitter {
 				throw new Error("Shadow git repo not initialized")
 			}
 
+			// Validate worktree configuration before proceeding
+			const worktree = await this.getShadowGitConfigWorktree(this.git)
+			if (worktree !== this.workspaceDir) {
+				throw new Error(
+					`Worktree mismatch detected. Expected: ${this.workspaceDir}, Found: ${worktree}. Aborting restore to prevent data loss.`,
+				)
+			}
+
+			// Verify the commit exists
+			try {
+				await this.git.catFile(["-t", commitHash])
+			} catch (error) {
+				throw new Error(`Invalid commit hash: ${commitHash}`)
+			}
+
 			const start = Date.now()
-			await this.git.clean("f", ["-d", "-f"])
+
+			// Simply reset to the target commit
+			// This is much safer than using git clean -d -f
 			await this.git.reset(["--hard", commitHash])
 
+			// Important: We do NOT use git clean here because it would remove untracked files
+			// git reset --hard only affects tracked files, leaving untracked files alone
+
 			// Remove all checkpoints after the specified commitHash.
 			const checkpointIndex = this._checkpoints.indexOf(commitHash)
 
@@ -322,6 +368,58 @@ export abstract class ShadowCheckpointService extends EventEmitter {
 		return super.once(event, listener)
 	}
 
+	/**
+	 * Recovery methods for data loss scenarios
+	 */
+
+	public async listStashes(): Promise<Array<{ index: number; message: string; date: string }>> {
+		if (!this.git) {
+			throw new Error("Shadow git repo not initialized")
+		}
+
+		try {
+			const stashList = await this.git.stashList()
+			return stashList.all.map((stash, index) => ({
+				index,
+				message: stash.message,
+				date: stash.date || new Date().toISOString(),
+			}))
+		} catch (error) {
+			const errorMessage = error instanceof Error ? error.message : String(error)
+			this.log(`[${this.constructor.name}#listStashes] failed to list stashes: ${errorMessage}`)
+			return []
+		}
+	}
+
+	public async recoverFromStash(stashIndex: number): Promise<void> {
+		if (!this.git) {
+			throw new Error("Shadow git repo not initialized")
+		}
+
+		try {
+			this.log(`[${this.constructor.name}#recoverFromStash] attempting to recover from stash@{${stashIndex}}`)
+
+			// Apply the stash without removing it from the stash list
+			// Use --index to also restore staged changes
+			await this.git.stash(["apply", "--index", `stash@{${stashIndex}}`])
+
+			this.log(`[${this.constructor.name}#recoverFromStash] successfully recovered from stash@{${stashIndex}}`)
+		} catch (error) {
+			// If applying with --index fails, try without it
+			try {
+				this.log(`[${this.constructor.name}#recoverFromStash] retrying without --index`)
+				await this.git.stash(["apply", `stash@{${stashIndex}}`])
+				this.log(
+					`[${this.constructor.name}#recoverFromStash] successfully recovered from stash@{${stashIndex}} (without --index)`,
+				)
+			} catch (retryError) {
+				const errorMessage = retryError instanceof Error ? retryError.message : String(retryError)
+				this.log(`[${this.constructor.name}#recoverFromStash] failed to recover from stash: ${errorMessage}`)
+				throw new Error(`Failed to recover from stash: ${errorMessage}`)
+			}
+		}
+	}
+
 	/**
 	 * Storage
 	 */

src/services/checkpoints/__tests__/ShadowCheckpointService.spec.ts

@@ -822,5 +822,137 @@ describe.each([[RepoPerTaskCheckpointService, "RepoPerTaskCheckpointService"]])(
 				expect(await fs.readFile(testFile, "utf-8")).toBe("Hello, world!")
 			})
 		})
+
+		describe(`${klass.name}#restoreCheckpoint safety`, () => {
+			it("does not use dangerous git clean command", async () => {
+				// This test verifies that we don't use git clean -d -f which would
+				// delete untracked files and potentially cause data loss
+
+				// Create initial checkpoint
+				await fs.writeFile(testFile, "Initial content")
+				const commit1 = await service.saveCheckpoint("Initial checkpoint")
+				expect(commit1?.commit).toBeTruthy()
+
+				// Make changes and create another checkpoint
+				await fs.writeFile(testFile, "Modified content")
+				const commit2 = await service.saveCheckpoint("Second checkpoint")
+				expect(commit2?.commit).toBeTruthy()
+
+				// Create a spy to monitor git commands
+				const gitSpy = vitest.spyOn(service["git"]!, "clean")
+
+				// Restore to first checkpoint
+				await service.restoreCheckpoint(commit1!.commit)
+
+				// Verify tracked file was restored
+				expect(await fs.readFile(testFile, "utf-8")).toBe("Initial content")
+
+				// Verify that git clean was NOT called
+				expect(gitSpy).not.toHaveBeenCalled()
+
+				gitSpy.mockRestore()
+			})
+
+			it("validates worktree configuration before restoration", async () => {
+				// Create a checkpoint
+				await fs.writeFile(testFile, "Test content")
+				const commit = await service.saveCheckpoint("Test checkpoint")
+				expect(commit?.commit).toBeTruthy()
+
+				// Create a new service instance with corrupted worktree
+				const corruptedShadowDir = path.join(tmpDir, `corrupted-${Date.now()}`)
+				const corruptedService = await klass.create({
+					taskId: "corrupted-test",
+					shadowDir: corruptedShadowDir,
+					workspaceDir: service.workspaceDir,
+					log: () => {},
+				})
+				await corruptedService.initShadowGit()
+
+				// Manually corrupt the worktree configuration by modifying the internal state
+				// This simulates a corrupted git config without actually setting an invalid path
+				corruptedService["shadowGitConfigWorktree"] = "/some/wrong/path"
+
+				// Attempt to restore should fail with worktree mismatch error
+				await expect(corruptedService.restoreCheckpoint(commit!.commit)).rejects.toThrow(
+					"Worktree mismatch detected",
+				)
+
+				// Using the original service (with correct worktree) should succeed
+				await service.restoreCheckpoint(commit!.commit)
+				expect(await fs.readFile(testFile, "utf-8")).toBe("Test content")
+
+				// Clean up
+				await fs.rm(corruptedShadowDir, { recursive: true, force: true })
+			})
+
+			it("validates commit hash before attempting restoration", async () => {
+				// Try to restore with an invalid commit hash
+				const invalidHash = "invalid-commit-hash-12345"
+
+				await expect(service.restoreCheckpoint(invalidHash)).rejects.toThrow(
+					"Invalid commit hash: " + invalidHash,
+				)
+			})
+
+			it("safely restores without data loss", async () => {
+				// Create initial checkpoint
+				await fs.writeFile(testFile, "Initial content")
+				const commit1 = await service.saveCheckpoint("Initial checkpoint")
+				expect(commit1?.commit).toBeTruthy()
+
+				// Make changes and save another checkpoint
+				await fs.writeFile(testFile, "Modified content")
+				const newFile = path.join(service.workspaceDir, "new-file.txt")
+				await fs.writeFile(newFile, "New file content")
+				const commit2 = await service.saveCheckpoint("Second checkpoint")
+				expect(commit2?.commit).toBeTruthy()
+
+				// Restore to initial checkpoint
+				await service.restoreCheckpoint(commit1!.commit)
+
+				// Verify restoration worked correctly
+				expect(await fs.readFile(testFile, "utf-8")).toBe("Initial content")
+
+				// The new file should not exist in the first checkpoint
+				await expect(fs.readFile(newFile, "utf-8")).rejects.toThrow()
+			})
+
+			it("handles restoration when no changes need to be stashed", async () => {
+				// Create and immediately restore a checkpoint (no changes to stash)
+				await fs.writeFile(testFile, "Test content")
+				const commit = await service.saveCheckpoint("Test checkpoint")
+				expect(commit?.commit).toBeTruthy()
+
+				// Restore immediately without making any changes
+				await expect(service.restoreCheckpoint(commit!.commit)).resolves.not.toThrow()
+				expect(await fs.readFile(testFile, "utf-8")).toBe("Test content")
+			})
+
+			it("properly removes files that should not exist in restored checkpoint", async () => {
+				// Create initial checkpoint
+				await fs.writeFile(testFile, "Initial content")
+				const commit1 = await service.saveCheckpoint("Initial checkpoint")
+				expect(commit1?.commit).toBeTruthy()
+
+				// Add a new tracked file
+				const trackedFile = path.join(service.workspaceDir, "tracked-file.txt")
+				await fs.writeFile(trackedFile, "This file will be tracked")
+				const commit2 = await service.saveCheckpoint("Added tracked file")
+				expect(commit2?.commit).toBeTruthy()
+
+				// Verify file exists
+				expect(await fs.readFile(trackedFile, "utf-8")).toBe("This file will be tracked")
+
+				// Restore to initial checkpoint (before the file existed)
+				await service.restoreCheckpoint(commit1!.commit)
+
+				// Verify the tracked file was removed
+				await expect(fs.readFile(trackedFile, "utf-8")).rejects.toThrow()
+
+				// Verify original file is still correct
+				expect(await fs.readFile(testFile, "utf-8")).toBe("Initial content")
+			})
+		})
 	},
 )