Edit/Delete with Checkpoints

src/core/checkpoints/__tests__/utils.test.ts

@@ -0,0 +1,107 @@
+import { describe, it, expect } from "vitest"
+import { isValidCheckpoint, hasValidCheckpoint, extractCheckpoint, type ValidCheckpoint } from "../utils"
+
+describe("checkpoint utils", () => {
+	describe("isValidCheckpoint", () => {
+		it("should return true for valid checkpoint", () => {
+			const checkpoint: ValidCheckpoint = { hash: "abc123" }
+			expect(isValidCheckpoint(checkpoint)).toBe(true)
+		})
+
+		it("should return false for null or undefined", () => {
+			expect(isValidCheckpoint(null)).toBe(false)
+			expect(isValidCheckpoint(undefined)).toBe(false)
+		})
+
+		it("should return false for non-object types", () => {
+			expect(isValidCheckpoint("string")).toBe(false)
+			expect(isValidCheckpoint(123)).toBe(false)
+			expect(isValidCheckpoint(true)).toBe(false)
+			expect(isValidCheckpoint([])).toBe(false)
+		})
+
+		it("should return false for objects without hash property", () => {
+			expect(isValidCheckpoint({})).toBe(false)
+			expect(isValidCheckpoint({ other: "property" })).toBe(false)
+		})
+
+		it("should return false for objects with non-string hash", () => {
+			expect(isValidCheckpoint({ hash: 123 })).toBe(false)
+			expect(isValidCheckpoint({ hash: null })).toBe(false)
+			expect(isValidCheckpoint({ hash: undefined })).toBe(false)
+			expect(isValidCheckpoint({ hash: {} })).toBe(false)
+			expect(isValidCheckpoint({ hash: [] })).toBe(false)
+		})
+
+		it("should return false for empty hash string", () => {
+			expect(isValidCheckpoint({ hash: "" })).toBe(false)
+		})
+
+		it("should return true for valid hash strings", () => {
+			expect(isValidCheckpoint({ hash: "a" })).toBe(true)
+			expect(isValidCheckpoint({ hash: "abc123def456" })).toBe(true)
+			expect(isValidCheckpoint({ hash: "commit-hash-with-dashes" })).toBe(true)
+		})
+	})
+
+	describe("hasValidCheckpoint", () => {
+		it("should return true for message with valid checkpoint", () => {
+			const message = { checkpoint: { hash: "abc123" } }
+			expect(hasValidCheckpoint(message)).toBe(true)
+		})
+
+		it("should return false for null or undefined message", () => {
+			expect(hasValidCheckpoint(null)).toBe(false)
+			expect(hasValidCheckpoint(undefined)).toBe(false)
+		})
+
+		it("should return false for non-object message", () => {
+			expect(hasValidCheckpoint("string")).toBe(false)
+			expect(hasValidCheckpoint(123)).toBe(false)
+		})
+
+		it("should return false for message without checkpoint property", () => {
+			expect(hasValidCheckpoint({})).toBe(false)
+			expect(hasValidCheckpoint({ text: "message" })).toBe(false)
+		})
+
+		it("should return false for message with invalid checkpoint", () => {
+			expect(hasValidCheckpoint({ checkpoint: null })).toBe(false)
+			expect(hasValidCheckpoint({ checkpoint: "invalid" })).toBe(false)
+			expect(hasValidCheckpoint({ checkpoint: {} })).toBe(false)
+			expect(hasValidCheckpoint({ checkpoint: { hash: "" } })).toBe(false)
+			expect(hasValidCheckpoint({ checkpoint: { hash: 123 } })).toBe(false)
+		})
+
+		it("should work as type guard", () => {
+			const message: unknown = { checkpoint: { hash: "abc123" }, other: "data" }
+			if (hasValidCheckpoint(message)) {
+				// TypeScript should know message has checkpoint property
+				expect(message.checkpoint.hash).toBe("abc123")
+			}
+		})
+	})
+
+	describe("extractCheckpoint", () => {
+		it("should extract valid checkpoint from message", () => {
+			const message = { checkpoint: { hash: "abc123" } }
+			const result = extractCheckpoint(message)
+			expect(result).toEqual({ hash: "abc123" })
+		})
+
+		it("should return undefined for message without valid checkpoint", () => {
+			expect(extractCheckpoint({})).toBeUndefined()
+			expect(extractCheckpoint({ checkpoint: null })).toBeUndefined()
+			expect(extractCheckpoint({ checkpoint: { hash: "" } })).toBeUndefined()
+			expect(extractCheckpoint(null)).toBeUndefined()
+			expect(extractCheckpoint(undefined)).toBeUndefined()
+		})
+
+		it("should return the same checkpoint object reference", () => {
+			const checkpoint = { hash: "abc123" }
+			const message = { checkpoint }
+			const result = extractCheckpoint(message)
+			expect(result).toBe(checkpoint)
+		})
+	})
+})

src/core/checkpoints/index.ts

@@ -14,6 +14,9 @@ import { DIFF_VIEW_URI_SCHEME } from "../../integrations/editor/DiffViewProvider
 
 import { CheckpointServiceOptions, RepoPerTaskCheckpointService } from "../../services/checkpoints"
 
+// Map to store pending checkpoint operations by taskId to prevent race conditions
+const pendingCheckpointOperations = new Map<string, Promise<any>>()
+
 export function getCheckpointService(cline: Task) {
 	if (!cline.enableCheckpoints) {
 		return undefined
@@ -150,35 +153,62 @@ async function getInitializedCheckpointService(
 }
 
 export async function checkpointSave(cline: Task, force = false) {
-	const service = getCheckpointService(cline)
+	const taskId = cline.taskId
 
-	if (!service) {
-		return
+	// Check if there's already a pending checkpoint operation for this task
+	const existingOperation = pendingCheckpointOperations.get(taskId)
+	if (existingOperation) {
+		// Return the existing Promise to prevent duplicate operations
+		return existingOperation
 	}
 
-	if (!service.isInitialized) {
-		const provider = cline.providerRef.deref()
-		provider?.log("[checkpointSave] checkpoints didn't initialize in time, disabling checkpoints for this task")
-		cline.enableCheckpoints = false
-		return
-	}
+	// Create a new checkpoint operation Promise
+	const checkpointOperation = (async () => {
+		try {
+			// Use getInitializedCheckpointService to wait for initialization
+			const service = await getInitializedCheckpointService(cline)
 
-	TelemetryService.instance.captureCheckpointCreated(cline.taskId)
+			if (!service) {
+				return
+			}
 
-	// Start the checkpoint process in the background.
-	return service.saveCheckpoint(`Task: ${cline.taskId}, Time: ${Date.now()}`, { allowEmpty: force }).catch((err) => {
-		console.error("[Task#checkpointSave] caught unexpected error, disabling checkpoints", err)
-		cline.enableCheckpoints = false
-	})
+			TelemetryService.instance.captureCheckpointCreated(cline.taskId)
+
+			// Start the checkpoint process in the background.
+			return await service.saveCheckpoint(`Task: ${cline.taskId}, Time: ${Date.now()}`, { allowEmpty: force })
+		} catch (err) {
+			console.error("[Task#checkpointSave] caught unexpected error, disabling checkpoints", err)
+			cline.enableCheckpoints = false
+			return undefined
+		}
+	})()
+
+	// Store the operation in the Map
+	pendingCheckpointOperations.set(taskId, checkpointOperation)
+
+	// Clean up the Map entry after the operation completes (success or failure)
+	checkpointOperation
+		.finally(() => {
+			pendingCheckpointOperations.delete(taskId)
+		})
+		.catch(() => {
+			// Error already handled above, this catch prevents unhandled rejection
+		})
+
+	return checkpointOperation
 }
 
 export type CheckpointRestoreOptions = {
 	ts: number
 	commitHash: string
 	mode: "preview" | "restore"
+	operation?: "delete" | "edit" // Optional to maintain backward compatibility
 }
 
-export async function checkpointRestore(cline: Task, { ts, commitHash, mode }: CheckpointRestoreOptions) {
+export async function checkpointRestore(
+	cline: Task,
+	{ ts, commitHash, mode, operation = "delete" }: CheckpointRestoreOptions,
+) {
 	const service = await getInitializedCheckpointService(cline)
 
 	if (!service) {
@@ -207,7 +237,10 @@ export async function checkpointRestore(cline: Task, { ts, commitHash, mode }: C
 				cline.combineMessages(deletedMessages),
 			)
 
-			await cline.overwriteClineMessages(cline.clineMessages.slice(0, index + 1))
+			// For delete operations, exclude the checkpoint message itself
+			// For edit operations, include the checkpoint message (to be edited)
+			const endIndex = operation === "edit" ? index + 1 : index
+			await cline.overwriteClineMessages(cline.clineMessages.slice(0, endIndex))
 
 			// TODO: Verify that this is working as expected.
 			await cline.say(

src/core/checkpoints/utils.ts

@@ -0,0 +1,51 @@
+/**
+ * Checkpoint-related utilities and type definitions
+ */
+
+/**
+ * Represents a valid checkpoint with required properties
+ */
+export interface ValidCheckpoint {
+	hash: string
+}
+
+/**
+ * Type guard to check if an object is a valid checkpoint
+ * @param checkpoint - The object to check
+ * @returns True if the checkpoint is valid, false otherwise
+ */
+export function isValidCheckpoint(checkpoint: unknown): checkpoint is ValidCheckpoint {
+	return (
+		checkpoint !== null &&
+		checkpoint !== undefined &&
+		typeof checkpoint === "object" &&
+		"hash" in checkpoint &&
+		typeof (checkpoint as any).hash === "string" &&
+		(checkpoint as any).hash.length > 0 // Ensure hash is not empty
+	)
+}
+
+/**
+ * Validates if a message has a valid checkpoint for restoration
+ * @param message - The message object to check
+ * @returns True if the message contains a valid checkpoint, false otherwise
+ */
+export function hasValidCheckpoint(message: unknown): message is { checkpoint: ValidCheckpoint } {
+	if (!message || typeof message !== "object" || !("checkpoint" in message)) {
+		return false
+	}
+
+	return isValidCheckpoint((message as any).checkpoint)
+}
+
+/**
+ * Extracts a valid checkpoint from a message if it exists
+ * @param message - The message object to extract from
+ * @returns The valid checkpoint or undefined
+ */
+export function extractCheckpoint(message: unknown): ValidCheckpoint | undefined {
+	if (hasValidCheckpoint(message)) {
+		return message.checkpoint
+	}
+	return undefined
+}

src/core/task/Task.ts

@@ -180,6 +180,7 @@ export class Task extends EventEmitter<ClineEvents> {
 	// LLM Messages & Chat Messages
 	apiConversationHistory: ApiMessage[] = []
 	clineMessages: ClineMessage[] = []
+	public pendingUserMessageCheckpoint?: Record<string, unknown>
 
 	// Ask
 	private askResponse?: ClineAskResponse
@@ -517,7 +518,14 @@ export class Task extends EventEmitter<ClineEvents> {
 			await this.addToClineMessages({ ts: askTs, type: "ask", ask: type, text, isProtected })
 		}
 
-		await pWaitFor(() => this.askResponse !== undefined || this.lastMessageTs !== askTs, { interval: 100 })
+		await pWaitFor(() => this.askResponse !== undefined || this.lastMessageTs !== askTs || this.abort, {
+			interval: 100,
+		})
+
+		if (this.abort) {
+			// Task was aborted, return a default response
+			return { response: "messageResponse", text: undefined, images: undefined }
+		}
 
 		if (this.lastMessageTs !== askTs) {
 			// Could happen if we send multiple asks in a row i.e. with
@@ -708,15 +716,31 @@ export class Task extends EventEmitter<ClineEvents> {
 				this.lastMessageTs = sayTs
 			}
 
-			await this.addToClineMessages({
-				ts: sayTs,
-				type: "say",
-				say: type,
-				text,
-				images,
-				checkpoint,
-				contextCondense,
-			})
+			if (type === "user_feedback") {
+				// Automatically use and clear the pending checkpoint for user_feedback messages
+				const feedbackCheckpoint = checkpoint || this.pendingUserMessageCheckpoint
+				this.pendingUserMessageCheckpoint = undefined // Clear it after use
+
+				await this.addToClineMessages({
+					ts: sayTs,
+					type: "say",
+					say: type,
+					text,
+					images,
+					checkpoint: feedbackCheckpoint,
+					contextCondense,
+				})
+			} else {
+				await this.addToClineMessages({
+					ts: sayTs,
+					type: "say",
+					say: type,
+					text,
+					images,
+					checkpoint,
+					contextCondense,
+				})
+			}
 		}
 	}
 
@@ -743,6 +767,7 @@ export class Task extends EventEmitter<ClineEvents> {
 		this.apiConversationHistory = []
 		await this.providerRef.deref()?.postStateToWebview()
 
+		// Checkpoint will be saved in handleWebviewAskResponse before this message is created
 		await this.say("text", task, images)
 		this.isInitialized = true
 
@@ -843,7 +868,6 @@ export class Task extends EventEmitter<ClineEvents> {
 			responseText = text
 			responseImages = images
 		}
-
 		// Make sure that the api conversation history can be resumed by the API,
 		// even if it goes out of sync with cline messages.
 		let existingApiConversationHistory: ApiMessage[] = await this.getSavedApiConversationHistory()
@@ -1081,6 +1105,13 @@ export class Task extends EventEmitter<ClineEvents> {
 			this.abandoned = true
 		}
 
+		// Resolve any pending ask operations to prevent "Current ask promise was ignored" errors
+		if (this.askResponse === undefined) {
+			this.askResponse = "messageResponse"
+			this.askResponseText = undefined
+			this.askResponseImages = undefined
+		}
+
 		this.abort = true
 		this.emit("taskAborted")