Edit/Delete with Checkpoints

src/core/checkpoints/index.ts

@@ -16,6 +16,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
@@ -192,35 +195,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) {
@@ -249,7 +279,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/webview/ClineProvider.ts

@@ -81,6 +81,16 @@ export type ClineProviderEvents = {
 	clineCreated: [cline: Task]
 }
 
+interface PendingEditOperation {
+	messageTs: number
+	editedContent: string
+	images?: string[]
+	messageIndex: number
+	apiConversationHistoryIndex: number
+	timeoutId: NodeJS.Timeout
+	createdAt: number
+}
+
 class OrganizationAllowListViolationError extends Error {
 	constructor(message: string) {
 		super(message)
@@ -109,6 +119,8 @@ export class ClineProvider
 	protected mcpHub?: McpHub // Change from private to protected
 	private marketplaceManager: MarketplaceManager
 	private mdmService?: MdmService
+	private pendingOperations: Map<string, PendingEditOperation> = new Map()
+	private static readonly PENDING_OPERATION_TIMEOUT_MS = 30000 // 30 seconds
 
 	public isViewLaunched = false
 	public settingsImportedAt?: number
@@ -242,6 +254,72 @@ export class ClineProvider
 		await this.removeClineFromStack()
 	}
 
+	// Pending Edit Operations Management
+
+	/**
+	 * Sets a pending edit operation with automatic timeout cleanup
+	 */
+	public setPendingEditOperation(
+		operationId: string,
+		editData: {
+			messageTs: number
+			editedContent: string
+			images?: string[]
+			messageIndex: number
+			apiConversationHistoryIndex: number
+		},
+	): void {
+		// Clear any existing operation with the same ID
+		this.clearPendingEditOperation(operationId)
+
+		// Create timeout for automatic cleanup
+		const timeoutId = setTimeout(() => {
+			this.clearPendingEditOperation(operationId)
+			this.log(`[setPendingEditOperation] Automatically cleared stale pending operation: ${operationId}`)
+		}, ClineProvider.PENDING_OPERATION_TIMEOUT_MS)
+
+		// Store the operation
+		this.pendingOperations.set(operationId, {
+			...editData,
+			timeoutId,
+			createdAt: Date.now(),
+		})
+
+		this.log(`[setPendingEditOperation] Set pending operation: ${operationId}`)
+	}
+
+	/**
+	 * Gets a pending edit operation by ID
+	 */
+	private getPendingEditOperation(operationId: string): PendingEditOperation | undefined {
+		return this.pendingOperations.get(operationId)
+	}
+
+	/**
+	 * Clears a specific pending edit operation
+	 */
+	private clearPendingEditOperation(operationId: string): boolean {
+		const operation = this.pendingOperations.get(operationId)
+		if (operation) {
+			clearTimeout(operation.timeoutId)
+			this.pendingOperations.delete(operationId)
+			this.log(`[clearPendingEditOperation] Cleared pending operation: ${operationId}`)
+			return true
+		}
+		return false
+	}
+
+	/**
+	 * Clears all pending edit operations
+	 */
+	private clearAllPendingEditOperations(): void {
+		for (const [operationId, operation] of this.pendingOperations) {
+			clearTimeout(operation.timeoutId)
+		}
+		this.pendingOperations.clear()
+		this.log(`[clearAllPendingEditOperations] Cleared all pending operations`)
+	}
+
 	/*
 	VSCode extensions use the disposable pattern to clean up resources when the sidebar/editor tab is closed by the user or system. This applies to event listening, commands, interacting with the UI, etc.
 	- https://vscode-docs.readthedocs.io/en/stable/extensions/patterns-and-principles/
@@ -261,6 +339,10 @@ export class ClineProvider
 		await this.removeClineFromStack()
 		this.log("Cleared task")
 
+		// Clear all pending edit operations to prevent memory leaks
+		this.clearAllPendingEditOperations()
+		this.log("Cleared pending operations")
+
 		if (this.view && "dispose" in this.view) {
 			this.view.dispose()
 			this.log("Disposed webview")
@@ -605,6 +687,50 @@ export class ClineProvider
 		this.log(
 			`[subtasks] ${cline.parentTask ? "child" : "parent"} task ${cline.taskId}.${cline.instanceId} instantiated`,
 		)
+
+		// Check if there's a pending edit after checkpoint restoration
+		const operationId = `task-${cline.taskId}`
+		const pendingEdit = this.getPendingEditOperation(operationId)
+		if (pendingEdit) {
+			this.clearPendingEditOperation(operationId) // Clear the pending edit
+
+			this.log(`[initClineWithHistoryItem] Processing pending edit after checkpoint restoration`)
+
+			// Process the pending edit after a short delay to ensure the task is fully initialized
+			setTimeout(async () => {
+				try {
+					// Find the message index in the restored state
+					const { messageIndex, apiConversationHistoryIndex } = (() => {
+						const messageIndex = cline.clineMessages.findIndex((msg) => msg.ts === pendingEdit.messageTs)
+						const apiConversationHistoryIndex = cline.apiConversationHistory.findIndex(
+							(msg) => msg.ts === pendingEdit.messageTs,
+						)
+						return { messageIndex, apiConversationHistoryIndex }
+					})()
+
+					if (messageIndex !== -1) {
+						// Remove the target message and all subsequent messages
+						await cline.overwriteClineMessages(cline.clineMessages.slice(0, messageIndex))
+
+						if (apiConversationHistoryIndex !== -1) {
+							await cline.overwriteApiConversationHistory(
+								cline.apiConversationHistory.slice(0, apiConversationHistoryIndex),
+							)
+						}
+
+						// Process the edited message
+						await cline.handleWebviewAskResponse(
+							"messageResponse",
+							pendingEdit.editedContent,
+							pendingEdit.images,
+						)
+					}
+				} catch (error) {
+					this.log(`[initClineWithHistoryItem] Error processing pending edit: ${error}`)
+				}
+			}, 100) // Small delay to ensure task is fully ready
+		}
+
 		return cline
 	}