Better error handling

src/core/controller/checkpoints/checkpointRestore.ts

@@ -23,7 +23,11 @@ export async function checkpointRestore(controller: Controller, request: Checkpo
 		})
 
 		// NOTE: cancelTask awaits abortTask, which awaits diffViewProvider.revertChanges, which reverts any edited files, allowing us to reset to a checkpoint rather than running into a state where the revertChanges function is called alongside or after the checkpoint reset
-		await controller.task?.restoreCheckpoint(request.number, request.restoreType as ClineCheckpointRestore, request.offset)
+		await controller.task?.checkpointManager.restoreCheckpoint(
+			request.number,
+			request.restoreType as ClineCheckpointRestore,
+			request.offset,
+		)
 	}
 	return Empty.create({})
 }

src/core/task/TaskState.ts

@@ -43,7 +43,7 @@ export class TaskState {
 	// Error tracking
 	consecutiveMistakeCount: number = 0
 	didAutomaticallyRetryFailedApiRequest = false
-	checkpointTrackerErrorMessage?: string
+	checkpointManagerErrorMessage?: string
 
 	// Task Initialization
 	isInitialized = false

src/core/task/ToolExecutor.ts

@@ -105,7 +105,7 @@ export class ToolExecutor {
 			images?: string[],
 			files?: string[],
 			partial?: boolean,
-		) => Promise<undefined>,
+		) => Promise<number | undefined>,
 		private ask: (
 			type: ClineAsk,
 			text?: string,
@@ -116,7 +116,7 @@ export class ToolExecutor {
 			images?: string[]
 			files?: string[]
 		}>,
-		private saveCheckpoint: (isAttemptCompletionMessage?: boolean) => Promise<void>,
+		private saveCheckpoint: (isAttemptCompletionMessage?: boolean, completionMessageTs?: number) => Promise<void>,
 		private sayAndCreateMissingParamError: (toolName: ToolUseName, paramName: string, relPath?: string) => Promise<any>,
 		private removeLastPartialMessageIfExistsWithType: (type: "ask" | "say", askOrSay: ClineAsk | ClineSay) => Promise<void>,
 		private executeCommandTool: (command: string) => Promise<[boolean, any]>,
@@ -2267,16 +2267,18 @@ export class ToolExecutor {
 								// await this.ask("command", this.removeClosingTag(block, "command", command), block.partial).catch(
 								// 	() => {},
 								// )
+
 							}
 						} else {
 							// no command, still outputting partial result
-							await this.say(
+							const completionMessageTs = await this.say(
 								"completion_result",
 								this.removeClosingTag(block, "result", result),
 								undefined,
 								undefined,
 								block.partial,
 							)
+							await this.saveCheckpoint(true, completionMessageTs)
 						}
 						break
 					} else {
@@ -2298,8 +2300,14 @@ export class ToolExecutor {
 						if (command) {
 							if (lastMessage && lastMessage.ask !== "command") {
 								// haven't sent a command message yet so first send completion_result then command
-								await this.say("completion_result", result, undefined, undefined, false)
-								await this.saveCheckpoint(true)
+								const completionMessageTs = await this.say(
+									"completion_result",
+									result,
+									undefined,
+									undefined,
+									false,
+								)
+								await this.saveCheckpoint(true, completionMessageTs)
 								await addNewChangesFlagToLastCompletionResultMessage()
 								telemetryService.captureTaskCompleted(this.ulid)
 
@@ -2331,8 +2339,8 @@ export class ToolExecutor {
 							// user didn't reject, but the command may have output
 							commandResult = execCommandResult
 						} else {
-							await this.say("completion_result", result, undefined, undefined, false)
-							await this.saveCheckpoint(true)
+							const completionMessageTs = await this.say("completion_result", result, undefined, undefined, false)
+							await this.saveCheckpoint(true, completionMessageTs)
 							await addNewChangesFlagToLastCompletionResultMessage()
 							telemetryService.captureTaskCompleted(this.ulid)
 

src/core/task/index.ts

@@ -110,7 +110,7 @@ export class Task {
 	browserSession: BrowserSession
 	contextManager: ContextManager
 	private diffViewProvider: DiffViewProvider
-	private checkpointTracker?: CheckpointTracker
+	public checkpointManager: TaskCheckpointManager
 	private clineIgnoreController: ClineIgnoreController
 	private toolExecutor: ToolExecutor
 
@@ -238,6 +238,35 @@ export class Task {
 			updateTaskHistory: this.updateTaskHistory,
 		})
 
+		// Initialize file context tracker
+		this.fileContextTracker = new FileContextTracker(context, this.taskId)
+		this.modelContextTracker = new ModelContextTracker(context, this.taskId)
+
+		// Initialize checkpoint manager
+		this.checkpointManager = createTaskCheckpointManager(
+			{
+				taskId: this.taskId,
+			},
+			{
+				enableCheckpoints: enableCheckpointsSetting,
+			},
+			{
+				context,
+				diffViewProvider: this.diffViewProvider,
+				messageStateHandler: this.messageStateHandler,
+				fileContextTracker: this.fileContextTracker,
+			},
+			{
+				updateTaskHistory: this.updateTaskHistory,
+				say: this.say.bind(this),
+				cancelTask: this.cancelTask,
+			},
+			{
+				conversationHistoryDeletedRange: this.taskState.conversationHistoryDeletedRange,
+				checkpointManagerErrorMessage: this.taskState.checkpointManagerErrorMessage,
+			},
+		)
+
 		// Initialize file context tracker
 		this.fileContextTracker = new FileContextTracker(controller, this.taskId)
 		this.modelContextTracker = new ModelContextTracker(controller.context, this.taskId)
@@ -404,310 +433,6 @@ export class Task {
 		}
 	}
 
-	async restoreCheckpoint(messageTs: number, restoreType: ClineCheckpointRestore, offset?: number) {
-		const clineMessages = this.messageStateHandler.getClineMessages()
-		const messageIndex = clineMessages.findIndex((m) => m.ts === messageTs) - (offset || 0)
-		// Find the last message before messageIndex that has a lastCheckpointHash
-		const lastHashIndex = findLastIndex(clineMessages.slice(0, messageIndex), (m) => m.lastCheckpointHash !== undefined)
-		const message = clineMessages[messageIndex]
-		const lastMessageWithHash = clineMessages[lastHashIndex]
-
-		if (!message) {
-			console.error("Message not found", clineMessages)
-			return
-		}
-
-		let didWorkspaceRestoreFail = false
-
-		switch (restoreType) {
-			case "task":
-				break
-			case "taskAndWorkspace":
-			case "workspace":
-				if (!this.enableCheckpoints) {
-					HostProvider.window.showMessage({
-						type: ShowMessageType.ERROR,
-						message: "Checkpoints are disabled in settings.",
-					})
-					didWorkspaceRestoreFail = true
-					break
-				}
-
-				if (!this.checkpointTracker && !this.taskState.checkpointTrackerErrorMessage) {
-					try {
-						this.checkpointTracker = await CheckpointTracker.create(
-							this.taskId,
-							this.controller.context.globalStorageUri.fsPath,
-							this.enableCheckpoints,
-						)
-						this.messageStateHandler.setCheckpointTracker(this.checkpointTracker)
-					} catch (error) {
-						const errorMessage = error instanceof Error ? error.message : "Unknown error"
-						console.error("Failed to initialize checkpoint tracker:", errorMessage)
-						this.taskState.checkpointTrackerErrorMessage = errorMessage
-						await this.postStateToWebview()
-						HostProvider.window.showMessage({
-							type: ShowMessageType.ERROR,
-							message: errorMessage,
-						})
-						didWorkspaceRestoreFail = true
-					}
-				}
-				if (message.lastCheckpointHash && this.checkpointTracker) {
-					try {
-						await this.checkpointTracker.resetHead(message.lastCheckpointHash)
-					} catch (error) {
-						const errorMessage = error instanceof Error ? error.message : "Unknown error"
-						HostProvider.window.showMessage({
-							type: ShowMessageType.ERROR,
-							message: "Failed to restore checkpoint: " + errorMessage,
-						})
-						didWorkspaceRestoreFail = true
-					}
-				} else if (offset && lastMessageWithHash.lastCheckpointHash && this.checkpointTracker) {
-					try {
-						await this.checkpointTracker.resetHead(lastMessageWithHash.lastCheckpointHash)
-					} catch (error) {
-						const errorMessage = error instanceof Error ? error.message : "Unknown error"
-						HostProvider.window.showMessage({
-							type: ShowMessageType.ERROR,
-							message: "Failed to restore offsetcheckpoint: " + errorMessage,
-						})
-						didWorkspaceRestoreFail = true
-					}
-				} else if (!offset && lastMessageWithHash.lastCheckpointHash && this.checkpointTracker) {
-					// Fallback: restore to most recent checkpoint when target message has no checkpoint hash
-					console.warn(`Message ${messageTs} has no checkpoint hash, falling back to previous checkpoint`)
-					try {
-						await this.checkpointTracker.resetHead(lastMessageWithHash.lastCheckpointHash)
-					} catch (error) {
-						const errorMessage = error instanceof Error ? error.message : "Unknown error"
-						HostProvider.window.showMessage({
-							type: ShowMessageType.ERROR,
-							message: "Failed to restore checkpoint: " + errorMessage,
-						})
-						didWorkspaceRestoreFail = true
-					}
-				} else {
-					HostProvider.window.showMessage({
-						type: ShowMessageType.ERROR,
-						message: "Failed to restore checkpoint",
-					})
-				}
-				break
-		}
-
-		if (!didWorkspaceRestoreFail) {
-			switch (restoreType) {
-				case "task":
-				case "taskAndWorkspace": {
-					this.taskState.conversationHistoryDeletedRange = message.conversationHistoryDeletedRange
-					const apiConversationHistory = this.messageStateHandler.getApiConversationHistory()
-					const newConversationHistory = apiConversationHistory.slice(0, (message.conversationHistoryIndex || 0) + 2) // +1 since this index corresponds to the last user message, and another +1 since slice end index is exclusive
-					await this.messageStateHandler.overwriteApiConversationHistory(newConversationHistory)
-
-					// update the context history state
-					await this.contextManager.truncateContextHistory(
-						message.ts,
-						await ensureTaskDirectoryExists(this.getContext(), this.taskId),
-					)
-
-					// aggregate deleted api reqs info so we don't lose costs/tokens
-					const clineMessages = this.messageStateHandler.getClineMessages()
-					const deletedMessages = clineMessages.slice(messageIndex + 1)
-					const deletedApiReqsMetrics = getApiMetrics(combineApiRequests(combineCommandSequences(deletedMessages)))
-
-					// Detect files edited after this message timestamp for file context warning
-					// Only needed for task-only restores when a user edits a message or restores the task context, but not the files.
-					if (restoreType === "task") {
-						const filesEditedAfterMessage = await this.fileContextTracker.detectFilesEditedAfterMessage(
-							messageTs,
-							deletedMessages,
-						)
-						if (filesEditedAfterMessage.length > 0) {
-							await this.fileContextTracker.storePendingFileContextWarning(filesEditedAfterMessage)
-						}
-					}
-
-					const newClineMessages = clineMessages.slice(0, messageIndex + 1)
-					await this.messageStateHandler.overwriteClineMessages(newClineMessages) // calls saveClineMessages which saves historyItem
-
-					await this.say(
-						"deleted_api_reqs",
-						JSON.stringify({
-							tokensIn: deletedApiReqsMetrics.totalTokensIn,
-							tokensOut: deletedApiReqsMetrics.totalTokensOut,
-							cacheWrites: deletedApiReqsMetrics.totalCacheWrites,
-							cacheReads: deletedApiReqsMetrics.totalCacheReads,
-							cost: deletedApiReqsMetrics.totalCost,
-						} satisfies ClineApiReqInfo),
-					)
-					break
-				}
-				case "workspace":
-					break
-			}
-
-			switch (restoreType) {
-				case "task":
-					HostProvider.window.showMessage({
-						type: ShowMessageType.INFORMATION,
-						message: "Task messages have been restored to the checkpoint",
-					})
-					break
-				case "workspace":
-					HostProvider.window.showMessage({
-						type: ShowMessageType.INFORMATION,
-						message: "Workspace files have been restored to the checkpoint",
-					})
-					break
-				case "taskAndWorkspace":
-					HostProvider.window.showMessage({
-						type: ShowMessageType.INFORMATION,
-						message: "Task and workspace have been restored to the checkpoint",
-					})
-					break
-			}
-
-			if (restoreType !== "task") {
-				// Set isCheckpointCheckedOut flag on the message
-				// Find all checkpoint messages before this one
-				const checkpointMessages = this.messageStateHandler
-					.getClineMessages()
-					.filter((m) => m.say === "checkpoint_created")
-				const currentMessageIndex = checkpointMessages.findIndex((m) => m.ts === messageTs)
-
-				// Set isCheckpointCheckedOut to false for all checkpoint messages
-				checkpointMessages.forEach((m, i) => {
-					m.isCheckpointCheckedOut = i === currentMessageIndex
-				})
-			}
-
-			await this.messageStateHandler.saveClineMessagesAndUpdateHistory()
-
-			this.cancelTask() // the task is already cancelled by the provider beforehand, but we need to re-init to get the updated messages
-		} else {
-			sendRelinquishControlEvent()
-		}
-	}
-
-	async presentMultifileDiff(messageTs: number, seeNewChangesSinceLastTaskCompletion: boolean) {
-		try {
-			if (!this.enableCheckpoints) {
-				HostProvider.window.showMessage({
-					type: ShowMessageType.INFORMATION,
-					message: "Checkpoints are disabled in settings. Cannot show diff.",
-				})
-				return
-			}
-			// TODO: handle if this is called from outside original workspace, in which case we need to
-			// show user error message we can't show diff outside of workspace?
-			if (!this.checkpointTracker && !this.taskState.checkpointTrackerErrorMessage) {
-				try {
-					this.checkpointTracker = await CheckpointTracker.create(
-						this.taskId,
-						this.controller.context.globalStorageUri.fsPath,
-						this.enableCheckpoints,
-					)
-					this.messageStateHandler.setCheckpointTracker(this.checkpointTracker)
-				} catch (error) {
-					console.error("Failed to initialize checkpoint tracker:", error)
-					const errorMessage = error instanceof Error ? error.message : "Unknown error"
-					this.taskState.checkpointTrackerErrorMessage = errorMessage
-					await this.postStateToWebview()
-					HostProvider.window.showMessage({
-						type: ShowMessageType.ERROR,
-						message: errorMessage,
-					})
-					return
-				}
-			}
-			if (!this.checkpointTracker) {
-				return
-			}
-
-			showChangedFilesDiff(
-				this.messageStateHandler,
-				this.checkpointTracker,
-				messageTs,
-				seeNewChangesSinceLastTaskCompletion,
-			)
-		} finally {
-			sendRelinquishControlEvent()
-		}
-	}
-
-	async doesLatestTaskCompletionHaveNewChanges() {
-		if (!this.enableCheckpoints) {
-			return false
-		}
-
-		const clineMessages = this.messageStateHandler.getClineMessages()
-		const messageIndex = findLastIndex(clineMessages, (m) => m.say === "completion_result")
-		const message = clineMessages[messageIndex]
-		if (!message) {
-			console.error("Completion message not found")
-			return false
-		}
-		const hash = message.lastCheckpointHash
-		if (!hash) {
-			console.error("No checkpoint hash found")
-			return false
-		}
-
-		if (this.enableCheckpoints && !this.checkpointTracker && !this.taskState.checkpointTrackerErrorMessage) {
-			try {
-				this.checkpointTracker = await CheckpointTracker.create(
-					this.taskId,
-					this.controller.context.globalStorageUri.fsPath,
-					this.enableCheckpoints,
-				)
-				this.messageStateHandler.setCheckpointTracker(this.checkpointTracker)
-			} catch (error) {
-				const errorMessage = error instanceof Error ? error.message : "Unknown error"
-				console.error("Failed to initialize checkpoint tracker:", errorMessage)
-				return false
-			}
-		}
-
-		// Get last task completed
-		const lastTaskCompletedMessage = findLast(
-			this.messageStateHandler.getClineMessages().slice(0, messageIndex),
-			(m) => m.say === "completion_result",
-		)
-
-		try {
-			// Get last task completed
-			const lastTaskCompletedMessageCheckpointHash = lastTaskCompletedMessage?.lastCheckpointHash // ask is only used to relinquish control, its the last say we care about
-			// if undefined, then we get diff from beginning of git
-			// if (!lastTaskCompletedMessage) {
-			// 	console.error("No previous task completion message found")
-			// 	return
-			// }
-			// This value *should* always exist
-			const firstCheckpointMessageCheckpointHash = this.messageStateHandler
-				.getClineMessages()
-				.find((m) => m.say === "checkpoint_created")?.lastCheckpointHash
-
-			const previousCheckpointHash = lastTaskCompletedMessageCheckpointHash || firstCheckpointMessageCheckpointHash // either use the diff between the first checkpoint and the task completion, or the diff between the latest two task completions
-
-			if (!previousCheckpointHash) {
-				return false
-			}
-
-			// Get count of changed files between current state and commit
-			const changedFilesCount = (await this.checkpointTracker?.getDiffCount(previousCheckpointHash, hash)) || 0
-			if (changedFilesCount > 0) {
-				return true
-			}
-		} catch (error) {
-			console.error("Failed to get diff set:", error)
-			return false
-		}
-
-		return false
-	}
-
 	// Communicate with webview
 
 	// partial has three valid states true (partial message), false (completion of partial message), undefined (individual complete message)
@@ -850,7 +575,13 @@ export class Task {
 		this.taskState.askResponseFiles = files
 	}
 
-	async say(type: ClineSay, text?: string, images?: string[], files?: string[], partial?: boolean): Promise<undefined> {
+	async say(
+		type: ClineSay,
+		text?: string,
+		images?: string[],
+		files?: string[],
+		partial?: boolean,
+	): Promise<number | undefined> {
 		if (this.taskState.abort) {
 			throw new Error("Cline instance aborted")
 		}
@@ -868,6 +599,7 @@ export class Task {
 					lastMessage.partial = partial
 					const protoMessage = convertClineMessageToProto(lastMessage)
 					await sendPartialMessageEvent(protoMessage)
+					return undefined
 				} else {
 					// this is a new partial message, so add it with partial state
 					const sayTs = Date.now()
@@ -882,6 +614,7 @@ export class Task {
 						partial,
 					})
 					await this.postStateToWebview()
+					return sayTs
 				}
 			} else {
 				// partial=false means its a complete version of a previously partial message
@@ -899,6 +632,7 @@ export class Task {
 					// await this.postStateToWebview()
 					const protoMessage = convertClineMessageToProto(lastMessage)
 					await sendPartialMessageEvent(protoMessage) // more performant than an entire postStateToWebview
+					return undefined
 				} else {
 					// this is a new partial=false message, so add it like normal
 					const sayTs = Date.now()
@@ -912,6 +646,7 @@ export class Task {
 						files,
 					})
 					await this.postStateToWebview()
+					return sayTs
 				}
 			}
 		} else {
@@ -927,6 +662,7 @@ export class Task {
 				files,
 			})
 			await this.postStateToWebview()
+			return sayTs
 		}
 	}
 
@@ -999,12 +735,6 @@ export class Task {
 			console.error("Failed to initialize ClineIgnoreController:", error)
 			// Optionally, inform the user or handle the error appropriately
 		}
-		// UPDATE: we don't need this anymore since most tasks are now created with checkpoints enabled
-		// right now we let users init checkpoints for old tasks, assuming they're continuing them from the same workspace (which we never tied to tasks, so no way for us to know if it's opened in the right workspace)
-		// const doesShadowGitExist = await CheckpointTracker.doesShadowGitExist(this.taskId, this.controllerRef.deref())
-		// if (!doesShadowGitExist) {
-		// 	this.checkpointTrackerErrorMessage = "Checkpoints are only available for new tasks"
-		// }
 
 		const savedClineMessages = await getSavedClineMessages(this.getContext(), this.taskId)
 
@@ -1229,145 +959,18 @@ export class Task {
 		}
 	}
 
-	// Checkpoints
+	// Checkpoints logic moved to checkpointManager
 
-	async saveCheckpoint(isAttemptCompletionMessage: boolean = false) {
-		if (
-			!this.enableCheckpoints ||
-			this.taskState.checkpointTrackerErrorMessage?.includes("Checkpoints initialization timed out.")
-		) {
-			// If checkpoints are disabled or previously encountered a timeout error, do nothing.
-			return
-		}
-		// Set isCheckpointCheckedOut to false for all checkpoint_created messages
-		this.messageStateHandler.getClineMessages().forEach((message) => {
-			if (message.say === "checkpoint_created") {
-				message.isCheckpointCheckedOut = false
-			}
-		})
+	async saveCheckpoint(isAttemptCompletionMessage: boolean = false, completionMessageTs?: number) {
+		await this.checkpointManager.saveCheckpoint(isAttemptCompletionMessage, completionMessageTs)
+	}
 
-		if (!isAttemptCompletionMessage) {
-			// ensure we aren't creating a duplicate checkpoint
-			const lastMessage = this.messageStateHandler.getClineMessages().at(-1)
-			if (lastMessage?.say === "checkpoint_created") {
-				return
-			}
+	async presentMultifileDiff(messageTs: number, seeNewChangesSinceLastTaskCompletion: boolean) {
+		await this.checkpointManager.presentMultifileDiff(messageTs, seeNewChangesSinceLastTaskCompletion)
+	}
 
-			// Initialize checkpoint tracker if it doesn't exist
-			if (!this.checkpointTracker && !this.taskState.checkpointTrackerErrorMessage) {
-				try {
-					this.checkpointTracker = await CheckpointTracker.create(
-						this.taskId,
-						this.controller.context.globalStorageUri.fsPath,
-						this.enableCheckpoints,
-					)
-				} catch (error) {
-					const errorMessage = error instanceof Error ? error.message : "Unknown error"
-					console.error("Failed to initialize checkpoint tracker:", errorMessage)
-					this.taskState.checkpointTrackerErrorMessage = errorMessage
-					await this.postStateToWebview()
-					return
-				}
-			}
-
-			// Create a checkpoint commit and update clineMessages with a commitHash
-			if (this.checkpointTracker) {
-				// We are letting this run in a non-blocking way so that the UI doesn't freeze when creating checkpoints.
-				// We show that a checkpoint is created in the chatview, then in the background run the git operation (which can take multiple seconds for large shadow git repos), and once that's been completed update the previous checkpoint message with the newly created hash to be associated with.
-				// NOTE: the attempt completion flow is different in that it requires the latest checkpoint hash to be present before determining if it can present the 'see new changes' button. In ToolExecutor, when we call saveCheckpoint(true), we must make sure that the checkpoint hash is present in the last completion_result message before returning, since it is always followed by a addNewChangesFlagToLastCompletionResultMessage(), which calls doesLatestTaskCompletionHaveNewChanges() that uses the latest message hash to determine if there any changes since the last attempt_completion checkpoint.
-				await this.say("checkpoint_created")
-				this.checkpointTracker.commit().then(async (commitHash) => {
-					if (commitHash) {
-						const lastCheckpointMessageIndex = findLastIndex(
-							this.messageStateHandler.getClineMessages(),
-							(m) => m.say === "checkpoint_created",
-						)
-						if (lastCheckpointMessageIndex !== -1) {
-							await this.messageStateHandler.updateClineMessage(lastCheckpointMessageIndex, {
-								lastCheckpointHash: commitHash,
-							})
-						}
-					}
-				})
-			} // silently fails for now
-
-			//
-		} else {
-			// attempt completion requires checkpoint to be sync so that we can present button after attempt_completion
-			// Check if checkpoint tracker exists, if not, create it. Skip if there was a previous checkpoints initialization timeout error.
-			if (
-				!this.checkpointTracker &&
-				!this.taskState.checkpointTrackerErrorMessage?.includes("Checkpoints initialization timed out.")
-			) {
-				try {
-					this.checkpointTracker = await CheckpointTracker.create(
-						this.taskId,
-						this.controller.context.globalStorageUri.fsPath,
-						this.enableCheckpoints,
-					)
-					this.messageStateHandler.setCheckpointTracker(this.checkpointTracker)
-				} catch (error) {
-					const errorMessage = error instanceof Error ? error.message : "Unknown error"
-					console.error("Failed to initialize checkpoint tracker for attempt completion:", errorMessage)
-					return
-				}
-			}
-
-			if (
-				this.checkpointTracker &&
-				!this.taskState.checkpointTrackerErrorMessage?.includes("Checkpoints initialization timed out.")
-			) {
-				const commitHash = await this.checkpointTracker.commit()
-
-				// For attempt_completion, find the last completion_result message and set its checkpoint hash. This will be used to present the 'see new changes' button
-				const lastCompletionResultMessage = findLast(
-					this.messageStateHandler.getClineMessages(),
-					(m) => m.say === "completion_result" || m.ask === "completion_result",
-				)
-				if (lastCompletionResultMessage) {
-					lastCompletionResultMessage.lastCheckpointHash = commitHash
-					await this.messageStateHandler.saveClineMessagesAndUpdateHistory()
-				}
-			} else {
-				console.error("Checkpoint tracker does not exist and could not be initialized for attempt completion")
-			}
-		}
-
-		// if (commitHash) {
-
-		// Previously we checkpointed every message, but this is excessive and unnecessary.
-		// // Start from the end and work backwards until we find a tool use or another message with a hash
-		// for (let i = this.clineMessages.length - 1; i >= 0; i--) {
-		// 	const message = this.clineMessages[i]
-		// 	if (message.lastCheckpointHash) {
-		// 		// Found a message with a hash, so we can stop
-		// 		break
-		// 	}
-		// 	// Update this message with a hash
-		// 	message.lastCheckpointHash = commitHash
-
-		// 	// We only care about adding the hash to the last tool use (we don't want to add this hash to every prior message ie for tasks pre-checkpoint)
-		// 	const isToolUse =
-		// 		message.say === "tool" ||
-		// 		message.ask === "tool" ||
-		// 		message.say === "command" ||
-		// 		message.ask === "command" ||
-		// 		message.say === "completion_result" ||
-		// 		message.ask === "completion_result" ||
-		// 		message.ask === "followup" ||
-		// 		message.say === "use_mcp_server" ||
-		// 		message.ask === "use_mcp_server" ||
-		// 		message.say === "browser_action" ||
-		// 		message.say === "browser_action_launch" ||
-		// 		message.ask === "browser_action_launch"
-
-		// 	if (isToolUse) {
-		// 		break
-		// 	}
-		// }
-		// // Save the updated messages
-		// await this.saveClineMessagesAndUpdateHistory()
-		// }
+	async doesLatestTaskCompletionHaveNewChanges(): Promise<boolean> {
+		return await this.checkpointManager.doesLatestTaskCompletionHaveNewChanges()
 	}
 
 	// Tools
@@ -2062,59 +1665,26 @@ export class Task {
 		if (
 			isFirstRequest &&
 			this.enableCheckpoints &&
-			!this.checkpointTracker &&
-			!this.taskState.checkpointTrackerErrorMessage
+			//!this.checkpointManager &&
+			!this.taskState.checkpointManagerErrorMessage
 		) {
 			try {
-				// Warning Timer - If checkpoints take a while to to initialize, show a warning message
-				let checkpointsWarningTimer: NodeJS.Timeout | null = null
-				let checkpointsWarningShown = false
-
-				checkpointsWarningTimer = setTimeout(async () => {
-					if (!checkpointsWarningShown) {
-						checkpointsWarningShown = true
-						this.taskState.checkpointTrackerErrorMessage =
-							"Checkpoints are taking longer than expected to initialize. Working in a large repository? Consider re-opening Cline in a project that uses git, or disabling checkpoints."
-						await this.postStateToWebview()
-					}
-				}, 7_000)
-
-				// Timeout - If checkpoints take too long to initialize, warn user and disable checkpoints for the task
-				this.checkpointTracker = await pTimeout(
-					CheckpointTracker.create(
-						this.taskId,
-						this.controller.context.globalStorageUri.fsPath,
-						this.enableCheckpoints,
-					),
-					{
-						milliseconds: 15_000,
-						message:
-							"Checkpoints taking too long to initialize. Consider re-opening Cline in a project that uses git, or disabling checkpoints.",
-					},
-				)
-				if (checkpointsWarningTimer) {
-					clearTimeout(checkpointsWarningTimer)
-					checkpointsWarningTimer = null
-				}
+				await pTimeout(this.checkpointManager.checkpointTrackerCheckAndInit(), {
+					milliseconds: 15_000,
+					message:
+						"Checkpoints taking too long to initialize. Consider re-opening Cline in a project that uses git, or disabling checkpoints.",
+				})
 			} catch (error) {
 				const errorMessage = error instanceof Error ? error.message : "Unknown error"
-				console.error("Failed to initialize checkpoint tracker:", errorMessage)
-
-				// If the error was a timeout, we disabled all checkpoint operations for the rest of the task
-				if (errorMessage.includes("Checkpoints taking too long to initialize")) {
-					this.taskState.checkpointTrackerErrorMessage =
-						"Checkpoints initialization timed out. Consider re-opening Cline in a project that uses git, or disabling checkpoints."
-					await this.postStateToWebview()
-				} else {
-					this.taskState.checkpointTrackerErrorMessage = errorMessage // will be displayed right away since we saveClineMessages next which posts state to webview
-				}
+				console.error("Failed to initialize checkpoint manager:", errorMessage)
+				this.taskState.checkpointManagerErrorMessage = errorMessage // will be displayed right away since we saveClineMessages next which posts state to webview
 			}
 		}
 
 		// Now, if it's the first request AND checkpoints are enabled AND tracker was successfully initialized,
 		// then say "checkpoint_created" and perform the commit.
-		if (isFirstRequest && this.enableCheckpoints && this.checkpointTracker) {
-			const commitHash = await this.checkpointTracker.commit() // Actual commit
+		if (isFirstRequest && this.enableCheckpoints && this.checkpointManager) {
+			const commitHash = await this.checkpointManager.commit() // Actual commit
 			await this.say("checkpoint_created") // Now this is conditional
 			const lastCheckpointMessageIndex = findLastIndex(
 				this.messageStateHandler.getClineMessages(),
@@ -2131,11 +1701,11 @@ export class Task {
 		} else if (
 			isFirstRequest &&
 			this.enableCheckpoints &&
-			!this.checkpointTracker &&
-			this.taskState.checkpointTrackerErrorMessage
+			!this.checkpointManager &&
+			this.taskState.checkpointManagerErrorMessage
 		) {
 			// Checkpoints are enabled, but tracker failed to initialize.
-			// checkpointTrackerErrorMessage is already set and will be part of the state.
+			// checkpointManagerErrorMessage is already set and will be part of the state.
 			// No explicit UI message here, error message will be in ExtensionState.
 		}
 

src/integrations/checkpoints/index.ts

@@ -0,0 +1,817 @@
+import * as vscode from "vscode"
+import { findLast, findLastIndex } from "@shared/array"
+import { ClineCheckpointRestore } from "@shared/WebviewMessage"
+import { ClineMessage, ClineApiReqInfo, ClineSay } from "@shared/ExtensionMessage"
+import { HistoryItem } from "@shared/HistoryItem"
+import CheckpointTracker from "@integrations/checkpoints/CheckpointTracker"
+import { DIFF_VIEW_URI_SCHEME, DiffViewProvider } from "@integrations/editor/DiffViewProvider"
+import { MessageStateHandler } from "../../core/task/message-state"
+import { ensureTaskDirectoryExists } from "@core/storage/disk"
+import { sendRelinquishControlEvent } from "@core/controller/ui/subscribeToRelinquishControl"
+import { ContextManager } from "@core/context/context-management/ContextManager"
+import { getApiMetrics } from "@shared/getApiMetrics"
+import { combineApiRequests } from "@shared/combineApiRequests"
+import { combineCommandSequences } from "@shared/combineCommandSequences"
+import { FileContextTracker } from "@core/context/context-tracking/FileContextTracker"
+
+// Type definitions for better code organization
+type SayFunction = (
+	type: ClineSay,
+	text?: string,
+	images?: string[],
+	files?: string[],
+	partial?: boolean,
+) => Promise<number | undefined>
+type UpdateTaskHistoryFunction = (historyItem: HistoryItem) => Promise<HistoryItem[]>
+
+interface CheckpointManagerTask {
+	readonly taskId: string
+}
+interface CheckpointManagerConfig {
+	readonly enableCheckpoints: boolean
+}
+interface CheckpointManagerServices {
+	readonly fileContextTracker: FileContextTracker
+	readonly diffViewProvider: DiffViewProvider
+	readonly messageStateHandler: MessageStateHandler
+	readonly context: vscode.ExtensionContext
+}
+interface CheckpointManagerCallbacks {
+	readonly updateTaskHistory: UpdateTaskHistoryFunction
+	readonly cancelTask: () => Promise<void>
+	readonly say: SayFunction
+}
+interface CheckpointManagerInternalState {
+	conversationHistoryDeletedRange?: [number, number]
+	checkpointTracker?: CheckpointTracker
+	checkpointManagerErrorMessage?: string
+	checkpointTrackerInitPromise?: Promise<CheckpointTracker | undefined>
+}
+
+interface CheckpointRestoreStateUpdate {
+	conversationHistoryDeletedRange?: [number, number]
+	checkpointManagerErrorMessage?: string
+}
+
+/**
+ * TaskCheckpointManager
+ *
+ * A dedicated service for managing all checkpoint-related operations within a task.
+ * Provides a clean separation of concerns from the main Task class while maintaining
+ * full access to necessary dependencies and state.
+ *
+ * Public API:
+ * - saveCheckpoint: Creates a new checkpoint of the current workspace state
+ * - restoreCheckpoint: Restores the task to a previous checkpoint
+ * - presentMultifileDiff: Displays a multi-file diff view between checkpoints
+ * - doesLatestTaskCompletionHaveNewChanges: Checks if the latest task completion has new changes, used by the "See New Changes" button
+ *
+ * This class is designed as the main interface between the task and the checkpoint system. It is responsible for:
+ * - Task-specific checkpoint operations (save/restore/diff)
+ * - State management and coordination with other Task components
+ * - Interaction with message state, file context tracking etc.
+ * - User interaction (error messages, notifications)
+ *
+ * For checkpoint operations, the CheckpointTracker class is used to interact with the underlying git logic.
+ */
+export class TaskCheckpointManager {
+	private readonly task: CheckpointManagerTask
+	private readonly config: CheckpointManagerConfig
+	private readonly services: CheckpointManagerServices
+	private readonly callbacks: CheckpointManagerCallbacks
+
+	private state: CheckpointManagerInternalState
+
+	constructor(
+		task: CheckpointManagerTask,
+		config: CheckpointManagerConfig,
+		services: CheckpointManagerServices,
+		callbacks: CheckpointManagerCallbacks,
+		initialState: CheckpointManagerInternalState,
+	) {
+		this.task = Object.freeze(task)
+		this.config = config
+		this.services = services
+		this.callbacks = Object.freeze(callbacks)
+		this.state = { ...initialState }
+	}
+
+	// ============================================================================
+	// Public API - Core checkpoints operations
+	// ============================================================================
+
+	/**
+	 * Creates a checkpoint of the current workspace state
+	 * @param isAttemptCompletionMessage - Whether this checkpoint is for an attempt completion message
+	 * @param completionMessageTs - Optional timestamp of the completion message to update with checkpoint hash
+	 */
+	async saveCheckpoint(isAttemptCompletionMessage: boolean = false, completionMessageTs?: number): Promise<void> {
+		try {
+			// If checkpoints are disabled, return early
+			if (!this.config.enableCheckpoints) {
+				return
+			}
+
+			// Set isCheckpointCheckedOut to false for all prior checkpoint_created messages
+			const clineMessages = this.services.messageStateHandler.getClineMessages()
+			clineMessages.forEach((message) => {
+				if (message.say === "checkpoint_created") {
+					message.isCheckpointCheckedOut = false
+				}
+			})
+
+			// Prevent repetitive checkpointTracker initialization errors on non-attempt completion messages
+			if (!this.state.checkpointTracker && !isAttemptCompletionMessage && !this.state.checkpointManagerErrorMessage) {
+				await this.checkpointTrackerCheckAndInit()
+			}
+			// attempt completion messages give it one last chance
+			else if (!this.state.checkpointTracker && isAttemptCompletionMessage) {
+				await this.checkpointTrackerCheckAndInit()
+			}
+
+			// Critical failure to initialize checkpoint tracker, return early
+			if (!this.state.checkpointTracker) {
+				console.error(
+					`[TaskCheckpointManager] Failed to save checkpoint for task ${this.task.taskId}: Checkpoint tracker not available`,
+				)
+				return
+			}
+
+			// Non attempt-completion messages call for a checkpoint_created message to be added
+			if (!isAttemptCompletionMessage) {
+				// Ensure we aren't creating back-to-back checkpoint_created messages
+				const lastMessage = clineMessages.at(-1)
+				if (lastMessage?.say === "checkpoint_created") {
+					return
+				}
+
+				// Create a new checkpoint_created message and asynchronously add the commitHash to the say message
+				const messageTs = await this.callbacks.say("checkpoint_created")
+				this.state.checkpointTracker
+					?.commit()
+					.then(async (commitHash) => {
+						if (messageTs) {
+							const targetMessage = this.services.messageStateHandler
+								.getClineMessages()
+								.find((m) => m.ts === messageTs)
+							if (targetMessage) {
+								targetMessage.lastCheckpointHash = commitHash
+								await this.services.messageStateHandler.saveClineMessagesAndUpdateHistory()
+							}
+						}
+					})
+					.catch((error) => {
+						console.error(
+							`[TaskCheckpointManager] Failed to create checkpoint commit for task ${this.task.taskId}:`,
+							error,
+						)
+					})
+			} else {
+				// attempt_completion messages are special
+				// First check last 3 messages to see if we already have a recent completion checkpoint
+				// If we do, skip creating a duplicate checkpoint
+				const lastFiveclineMessages = this.services.messageStateHandler.getClineMessages().slice(-3)
+				const lastCompletionResultMessage = findLast(lastFiveclineMessages, (m) => m.say === "completion_result")
+				if (lastCompletionResultMessage?.lastCheckpointHash) {
+					console.log("Completion checkpoint already exists, skipping duplicate checkpoint creation")
+					return
+				}
+
+				// For attempt_completion, commit then update the completion_result message with the checkpoint hash
+				if (this.state.checkpointTracker) {
+					const commitHash = await this.state.checkpointTracker.commit()
+
+					// If a completionMessageTs is provided, update that specific message with the checkpoint hash
+					if (completionMessageTs) {
+						const targetMessage = this.services.messageStateHandler
+							.getClineMessages()
+							.find((m) => m.ts === completionMessageTs)
+						if (targetMessage) {
+							targetMessage.lastCheckpointHash = commitHash
+							await this.services.messageStateHandler.saveClineMessagesAndUpdateHistory()
+						}
+					} else {
+						// Fallback to findLast if no timestamp provided - update the last completion_result message
+						if (lastCompletionResultMessage) {
+							lastCompletionResultMessage.lastCheckpointHash = commitHash
+							await this.services.messageStateHandler.saveClineMessagesAndUpdateHistory()
+						}
+					}
+				} else {
+					console.error(
+						`[TaskCheckpointManager] Checkpoint tracker does not exist and could not be initialized for attempt completion for task ${this.task.taskId}`,
+					)
+				}
+			}
+		} catch (error) {
+			const errorMessage = error instanceof Error ? error.message : "Unknown error"
+			console.error(`[TaskCheckpointManager] Failed to save checkpoint for task ${this.task.taskId}:`, errorMessage)
+		}
+	}
+
+	/**
+	 * Restores a checkpoint by message timestamp
+	 * @param messageTs - Timestamp of the message to restore to
+	 * @param restoreType - Type of restoration (task, workspace, or both)
+	 * @param offset - Optional offset for the message index
+	 * @returns checkpointManagerStateUpdate with any state changes that need to be applied
+	 */
+	async restoreCheckpoint(
+		messageTs: number,
+		restoreType: ClineCheckpointRestore,
+		offset?: number,
+	): Promise<CheckpointRestoreStateUpdate> {
+		try {
+			const clineMessages = this.services.messageStateHandler.getClineMessages()
+			const messageIndex = clineMessages.findIndex((m) => m.ts === messageTs) - (offset || 0)
+			// Find the last message before messageIndex that has a lastCheckpointHash
+			const lastHashIndex = findLastIndex(clineMessages.slice(0, messageIndex), (m) => m.lastCheckpointHash !== undefined)
+			const message = clineMessages[messageIndex]
+			const lastMessageWithHash = clineMessages[lastHashIndex]
+
+			if (!message) {
+				console.error(`[TaskCheckpointManager] Message not found for timestamp ${messageTs} in task ${this.task.taskId}`)
+				return {}
+			}
+
+			let didWorkspaceRestoreFail = false
+
+			switch (restoreType) {
+				case "task":
+					break
+				case "taskAndWorkspace":
+				case "workspace":
+					if (!this.config.enableCheckpoints) {
+						const errorMessage = "Checkpoints are disabled in settings."
+						console.error(`[TaskCheckpointManager] ${errorMessage} for task ${this.task.taskId}`)
+						vscode.window.showErrorMessage(errorMessage)
+						didWorkspaceRestoreFail = true
+						break
+					}
+
+					if (!this.state.checkpointTracker && !this.state.checkpointManagerErrorMessage) {
+						try {
+							this.state.checkpointTracker = await CheckpointTracker.create(
+								this.task.taskId,
+								this.services.context.globalStorageUri.fsPath,
+								this.config.enableCheckpoints,
+							)
+							this.services.messageStateHandler.setCheckpointTracker(this.state.checkpointTracker)
+						} catch (error) {
+							const errorMessage = error instanceof Error ? error.message : "Unknown error"
+							console.error(
+								`[TaskCheckpointManager] Failed to initialize checkpoint tracker for task ${this.task.taskId}:`,
+								errorMessage,
+							)
+							this.state.checkpointManagerErrorMessage = errorMessage
+							vscode.window.showErrorMessage(errorMessage)
+							didWorkspaceRestoreFail = true
+						}
+					}
+					if (message.lastCheckpointHash && this.state.checkpointTracker) {
+						try {
+							await this.state.checkpointTracker.resetHead(message.lastCheckpointHash)
+						} catch (error) {
+							const errorMessage = error instanceof Error ? error.message : "Unknown error"
+							console.error(
+								`[TaskCheckpointManager] Failed to restore checkpoint for task ${this.task.taskId}:`,
+								errorMessage,
+							)
+							vscode.window.showErrorMessage("Failed to restore checkpoint: " + errorMessage)
+							didWorkspaceRestoreFail = true
+						}
+					} else if (offset && lastMessageWithHash.lastCheckpointHash && this.state.checkpointTracker) {
+						try {
+							await this.state.checkpointTracker.resetHead(lastMessageWithHash.lastCheckpointHash)
+						} catch (error) {
+							const errorMessage = error instanceof Error ? error.message : "Unknown error"
+							console.error(
+								`[TaskCheckpointManager] Failed to restore offset checkpoint for task ${this.task.taskId}:`,
+								errorMessage,
+							)
+							vscode.window.showErrorMessage("Failed to restore offsetcheckpoint: " + errorMessage)
+							didWorkspaceRestoreFail = true
+						}
+					} else if (!offset && lastMessageWithHash.lastCheckpointHash && this.state.checkpointTracker) {
+						// Fallback: restore to most recent checkpoint when target message has no checkpoint hash
+						console.warn(
+							`[TaskCheckpointManager] Message ${messageTs} has no checkpoint hash, falling back to previous checkpoint for task ${this.task.taskId}`,
+						)
+						try {
+							await this.state.checkpointTracker.resetHead(lastMessageWithHash.lastCheckpointHash)
+						} catch (error) {
+							const errorMessage = error instanceof Error ? error.message : "Unknown error"
+							console.error(
+								`[TaskCheckpointManager] Failed to restore fallback checkpoint for task ${this.task.taskId}:`,
+								errorMessage,
+							)
+							vscode.window.showErrorMessage("Failed to restore checkpoint: " + errorMessage)
+							didWorkspaceRestoreFail = true
+						}
+					} else {
+						const errorMessage = "Failed to restore checkpoint: No valid checkpoint hash found"
+						console.error(`[TaskCheckpointManager] ${errorMessage} for task ${this.task.taskId}`)
+						vscode.window.showErrorMessage(errorMessage)
+						didWorkspaceRestoreFail = true
+					}
+					break
+			}
+
+			const checkpointManagerStateUpdate: CheckpointRestoreStateUpdate = {}
+
+			if (!didWorkspaceRestoreFail) {
+				await this.handleSuccessfulRestore(restoreType, message, messageIndex, messageTs)
+
+				// Collect state updates
+				if (this.state.conversationHistoryDeletedRange !== undefined) {
+					checkpointManagerStateUpdate.conversationHistoryDeletedRange = this.state.conversationHistoryDeletedRange
+				}
+			} else {
+				sendRelinquishControlEvent()
+
+				if (this.state.checkpointManagerErrorMessage !== undefined) {
+					checkpointManagerStateUpdate.checkpointManagerErrorMessage = this.state.checkpointManagerErrorMessage
+				}
+			}
+
+			return checkpointManagerStateUpdate
+		} catch (error) {
+			const errorMessage = error instanceof Error ? error.message : "Unknown error"
+			console.error(`[TaskCheckpointManager] Failed to restore checkpoint for task ${this.task.taskId}:`, errorMessage)
+			sendRelinquishControlEvent()
+			return {
+				checkpointManagerErrorMessage: errorMessage,
+			}
+		}
+	}
+
+	/**
+	 * Presents a multi-file diff view between checkpoints
+	 * @param messageTs - Timestamp of the message to show diff for
+	 * @param seeNewChangesSinceLastTaskCompletion - Whether to show changes since last completion
+	 */
+	async presentMultifileDiff(messageTs: number, seeNewChangesSinceLastTaskCompletion: boolean): Promise<void> {
+		const relinquishButton = () => {
+			sendRelinquishControlEvent()
+		}
+
+		try {
+			if (!this.config.enableCheckpoints) {
+				const errorMessage = "Checkpoints are disabled in settings. Cannot show diff."
+				console.error(`[TaskCheckpointManager] ${errorMessage} for task ${this.task.taskId}`)
+				vscode.window.showInformationMessage(errorMessage)
+				relinquishButton()
+				return
+			}
+
+			console.log(`[TaskCheckpointManager] presentMultifileDiff for task ${this.task.taskId}, messageTs: ${messageTs}`)
+			const clineMessages = this.services.messageStateHandler.getClineMessages()
+			const messageIndex = clineMessages.findIndex((m) => m.ts === messageTs)
+			const message = clineMessages[messageIndex]
+			if (!message) {
+				console.error(`[TaskCheckpointManager] Message not found for timestamp ${messageTs} in task ${this.task.taskId}`)
+				relinquishButton()
+				return
+			}
+			const hash = message.lastCheckpointHash
+			if (!hash) {
+				console.error(
+					`[TaskCheckpointManager] No checkpoint hash found for message ${messageTs} in task ${this.task.taskId}`,
+				)
+				relinquishButton()
+				return
+			}
+
+			// Initialize checkpoint tracker if needed
+			if (!this.state.checkpointTracker && this.config.enableCheckpoints && !this.state.checkpointManagerErrorMessage) {
+				try {
+					this.state.checkpointTracker = await CheckpointTracker.create(
+						this.task.taskId,
+						this.services.context.globalStorageUri.fsPath,
+						this.config.enableCheckpoints,
+					)
+					this.services.messageStateHandler.setCheckpointTracker(this.state.checkpointTracker)
+				} catch (error) {
+					const errorMessage = error instanceof Error ? error.message : "Unknown error"
+					console.error(
+						`[TaskCheckpointManager] Failed to initialize checkpoint tracker for task ${this.task.taskId}:`,
+						errorMessage,
+					)
+					this.state.checkpointManagerErrorMessage = errorMessage
+					vscode.window.showErrorMessage(errorMessage)
+					relinquishButton()
+					return
+				}
+			}
+
+			if (!this.state.checkpointTracker) {
+				console.error(`[TaskCheckpointManager] Checkpoint tracker not available for task ${this.task.taskId}`)
+				vscode.window.showErrorMessage("Checkpoint tracker not available")
+				relinquishButton()
+				return
+			}
+
+			let changedFiles:
+				| {
+						relativePath: string
+						absolutePath: string
+						before: string
+						after: string
+				  }[]
+				| undefined
+
+			if (seeNewChangesSinceLastTaskCompletion) {
+				// Get last task completed
+				const lastTaskCompletedMessageCheckpointHash = findLast(
+					this.services.messageStateHandler.getClineMessages().slice(0, messageIndex),
+					(m) => m.say === "completion_result",
+				)?.lastCheckpointHash
+
+				// This value *should* always exist
+				const firstCheckpointMessageCheckpointHash = this.services.messageStateHandler
+					.getClineMessages()
+					.find((m) => m.say === "checkpoint_created")?.lastCheckpointHash
+
+				const previousCheckpointHash = lastTaskCompletedMessageCheckpointHash || firstCheckpointMessageCheckpointHash
+
+				if (!previousCheckpointHash) {
+					const errorMessage = "Unexpected error: No checkpoint hash found"
+					console.error(`[TaskCheckpointManager] ${errorMessage} for task ${this.task.taskId}`)
+					vscode.window.showErrorMessage(errorMessage)
+					relinquishButton()
+					return
+				}
+
+				// Get changed files between current state and commit
+				changedFiles = await this.state.checkpointTracker.getDiffSet(previousCheckpointHash, hash)
+				if (!changedFiles?.length) {
+					vscode.window.showInformationMessage("No changes found")
+					relinquishButton()
+					return
+				}
+			} else {
+				// Get changed files between current state and commit
+				changedFiles = await this.state.checkpointTracker.getDiffSet(hash)
+				if (!changedFiles?.length) {
+					vscode.window.showInformationMessage("No changes found")
+					relinquishButton()
+					return
+				}
+			}
+
+			// Open multi-diff editor
+			await vscode.commands.executeCommand(
+				"vscode.changes",
+				seeNewChangesSinceLastTaskCompletion ? "New changes" : "Changes since snapshot",
+				changedFiles.map((file) => [
+					vscode.Uri.file(file.absolutePath),
+					vscode.Uri.parse(`${DIFF_VIEW_URI_SCHEME}:${file.relativePath}`).with({
+						query: Buffer.from(file.before ?? "").toString("base64"),
+					}),
+					vscode.Uri.parse(`${DIFF_VIEW_URI_SCHEME}:${file.relativePath}`).with({
+						query: Buffer.from(file.after ?? "").toString("base64"),
+					}),
+				]),
+			)
+			relinquishButton()
+		} catch (error) {
+			const errorMessage = error instanceof Error ? error.message : "Unknown error"
+			console.error(`[TaskCheckpointManager] Failed to present multifile diff for task ${this.task.taskId}:`, errorMessage)
+			vscode.window.showErrorMessage("Failed to retrieve diff set: " + errorMessage)
+			relinquishButton()
+		}
+	}
+
+	/**
+	 * Creates a checkpoint commit in the underlying tracker
+	 * @returns Promise<string | undefined> The created commit hash, or undefined if failed
+	 */
+	async commit(): Promise<string | undefined> {
+		try {
+			if (!this.config.enableCheckpoints) {
+				return undefined
+			}
+
+			if (!this.state.checkpointTracker) {
+				await this.checkpointTrackerCheckAndInit()
+			}
+
+			if (!this.state.checkpointTracker) {
+				console.error(`[TaskCheckpointManager] Checkpoint tracker not available for commit in task ${this.task.taskId}`)
+				return undefined
+			}
+
+			return await this.state.checkpointTracker.commit()
+		} catch (error) {
+			const errorMessage = error instanceof Error ? error.message : "Unknown error"
+			console.error(
+				`[TaskCheckpointManager] Failed to create checkpoint commit for task ${this.task.taskId}:`,
+				errorMessage,
+			)
+			return undefined
+		}
+	}
+
+	/**
+	 * Checks if the latest task completion has new changes
+	 * @returns Promise<boolean> - True if there are new changes since last completion
+	 */
+	async doesLatestTaskCompletionHaveNewChanges(): Promise<boolean> {
+		try {
+			if (!this.config.enableCheckpoints) {
+				return false
+			}
+
+			const clineMessages = this.services.messageStateHandler.getClineMessages()
+			const messageIndex = findLastIndex(clineMessages, (m) => m.say === "completion_result")
+			const message = clineMessages[messageIndex]
+			if (!message) {
+				console.error(`[TaskCheckpointManager] Completion message not found for task ${this.task.taskId}`)
+				return false
+			}
+			const hash = message.lastCheckpointHash
+			if (!hash) {
+				console.error(
+					`[TaskCheckpointManager] No checkpoint hash found for completion message in task ${this.task.taskId}`,
+				)
+				return false
+			}
+
+			if (this.config.enableCheckpoints && !this.state.checkpointTracker && !this.state.checkpointManagerErrorMessage) {
+				try {
+					this.state.checkpointTracker = await CheckpointTracker.create(
+						this.task.taskId,
+						this.services.context.globalStorageUri.fsPath,
+						this.config.enableCheckpoints,
+					)
+					this.services.messageStateHandler.setCheckpointTracker(this.state.checkpointTracker)
+				} catch (error) {
+					const errorMessage = error instanceof Error ? error.message : "Unknown error"
+					console.error(
+						`[TaskCheckpointManager] Failed to initialize checkpoint tracker for task ${this.task.taskId}:`,
+						errorMessage,
+					)
+					return false
+				}
+			}
+
+			if (!this.state.checkpointTracker) {
+				console.error(`[TaskCheckpointManager] Checkpoint tracker not available for task ${this.task.taskId}`)
+				return false
+			}
+
+			// Get last task completed
+			const lastTaskCompletedMessage = findLast(
+				this.services.messageStateHandler.getClineMessages().slice(0, messageIndex),
+				(m) => m.say === "completion_result",
+			)
+
+			// Get last task completed
+			const lastTaskCompletedMessageCheckpointHash = lastTaskCompletedMessage?.lastCheckpointHash
+
+			// This value *should* always exist
+			const firstCheckpointMessageCheckpointHash = this.services.messageStateHandler
+				.getClineMessages()
+				.find((m) => m.say === "checkpoint_created")?.lastCheckpointHash
+
+			const previousCheckpointHash = lastTaskCompletedMessageCheckpointHash || firstCheckpointMessageCheckpointHash
+
+			if (!previousCheckpointHash) {
+				console.error(`[TaskCheckpointManager] No previous checkpoint hash found for task ${this.task.taskId}`)
+				return false
+			}
+
+			// Get count of changed files between current state and commit
+			const changedFilesCount = (await this.state.checkpointTracker.getDiffCount(previousCheckpointHash, hash)) || 0
+			return changedFilesCount > 0
+		} catch (error) {
+			const errorMessage = error instanceof Error ? error.message : "Unknown error"
+			console.error(`[TaskCheckpointManager] Failed to check for new changes in task ${this.task.taskId}:`, errorMessage)
+			return false
+		}
+	}
+
+	/**
+	 * Handles the successful restoration logic for different restore types
+	 */
+	// Largely unchanged from original Task class implementation
+	private async handleSuccessfulRestore(
+		restoreType: ClineCheckpointRestore,
+		message: ClineMessage,
+		messageIndex: number,
+		messageTs: number,
+	): Promise<void> {
+		switch (restoreType) {
+			case "task":
+			case "taskAndWorkspace":
+				// Update conversation history deleted range in our state
+				this.state.conversationHistoryDeletedRange = message.conversationHistoryDeletedRange
+
+				const apiConversationHistory = this.services.messageStateHandler.getApiConversationHistory()
+				const newConversationHistory = apiConversationHistory.slice(0, (message.conversationHistoryIndex || 0) + 2) // +1 since this index corresponds to the last user message, and another +1 since slice end index is exclusive
+				await this.services.messageStateHandler.overwriteApiConversationHistory(newConversationHistory)
+
+				// update the context history state
+				const contextManager = new ContextManager()
+				await contextManager.truncateContextHistory(
+					message.ts,
+					await ensureTaskDirectoryExists(this.getContext(), this.task.taskId),
+				)
+
+				// aggregate deleted api reqs info so we don't lose costs/tokens
+				const clineMessages = this.services.messageStateHandler.getClineMessages()
+				const deletedMessages = clineMessages.slice(messageIndex + 1)
+				const deletedApiReqsMetrics = getApiMetrics(combineApiRequests(combineCommandSequences(deletedMessages)))
+
+				// Detect files edited after this message timestamp for file context warning
+				// Only needed for task-only restores when a user edits a message or restores the task context, but not the files.
+				if (restoreType === "task") {
+					const filesEditedAfterMessage = await this.services.fileContextTracker.detectFilesEditedAfterMessage(
+						messageTs,
+						deletedMessages,
+					)
+					if (filesEditedAfterMessage.length > 0) {
+						await this.services.fileContextTracker.storePendingFileContextWarning(filesEditedAfterMessage)
+					}
+				}
+
+				const newClineMessages = clineMessages.slice(0, messageIndex + 1)
+				await this.services.messageStateHandler.overwriteClineMessages(newClineMessages) // calls saveClineMessages which saves historyItem
+
+				await this.callbacks.say(
+					"deleted_api_reqs",
+					JSON.stringify({
+						tokensIn: deletedApiReqsMetrics.totalTokensIn,
+						tokensOut: deletedApiReqsMetrics.totalTokensOut,
+						cacheWrites: deletedApiReqsMetrics.totalCacheWrites,
+						cacheReads: deletedApiReqsMetrics.totalCacheReads,
+						cost: deletedApiReqsMetrics.totalCost,
+					} satisfies ClineApiReqInfo),
+				)
+				break
+			case "workspace":
+				break
+		}
+
+		switch (restoreType) {
+			case "task":
+				vscode.window.showInformationMessage("Task messages have been restored to the checkpoint")
+				break
+			case "workspace":
+				vscode.window.showInformationMessage("Workspace files have been restored to the checkpoint")
+				break
+			case "taskAndWorkspace":
+				vscode.window.showInformationMessage("Task and workspace have been restored to the checkpoint")
+				break
+		}
+
+		if (restoreType !== "task") {
+			// Set isCheckpointCheckedOut flag on the message
+			// Find all checkpoint messages before this one
+			const checkpointMessages = this.services.messageStateHandler
+				.getClineMessages()
+				.filter((m) => m.say === "checkpoint_created")
+			const currentMessageIndex = checkpointMessages.findIndex((m) => m.ts === messageTs)
+
+			// Set isCheckpointCheckedOut to false for all checkpoint messages
+			checkpointMessages.forEach((m, i) => {
+				m.isCheckpointCheckedOut = i === currentMessageIndex
+			})
+		}
+
+		await this.services.messageStateHandler.saveClineMessagesAndUpdateHistory()
+
+		// Cancel and reinitialize the task to get updated messages
+		await this.callbacks.cancelTask()
+	}
+
+	// ============================================================================
+	// State management - interfaces for updating internal state
+	// ============================================================================
+
+	/**
+	 * Checks for an active checkpoint tracker instance, creates if needed
+	 * Uses promise-based synchronization to prevent race conditions when called concurrently
+	 */
+	async checkpointTrackerCheckAndInit(): Promise<CheckpointTracker | undefined> {
+		// If tracker already exists or there was an error, return immediately
+		if (this.state.checkpointTracker) {
+			return this.state.checkpointTracker
+		}
+
+		// If initialization is already in progress, wait for it to complete
+		if (this.state.checkpointTrackerInitPromise) {
+			return await this.state.checkpointTrackerInitPromise
+		}
+
+		// Start initialization and store the promise to prevent concurrent attempts
+		this.state.checkpointTrackerInitPromise = this.initializeCheckpointTracker()
+
+		try {
+			const tracker = await this.state.checkpointTrackerInitPromise
+			return tracker
+		} finally {
+			// Clear the promise once initialization is complete (success or failure)
+			this.state.checkpointTrackerInitPromise = undefined
+		}
+	}
+
+	/**
+	 * Internal method to actually create the checkpoint tracker
+	 */
+	private async initializeCheckpointTracker(): Promise<CheckpointTracker | undefined> {
+		try {
+			const tracker = await CheckpointTracker.create(
+				this.task.taskId,
+				this.services.context.globalStorageUri.fsPath,
+				this.config.enableCheckpoints,
+			)
+
+			// Update the state with the created tracker
+			this.state.checkpointTracker = tracker
+			return tracker
+		} catch (error) {
+			// Store error message to prevent future repetative initialization attempts
+			const errorMessage = error instanceof Error ? error.message : "Unknown error"
+			this.setcheckpointManagerErrorMessage(errorMessage)
+			console.error("Failed to initialize checkpoint tracker:", errorMessage)
+			// TODO - Do we need to post state to webview here? TBD
+			return undefined
+		}
+	}
+
+	/**
+	 * Updates the checkpoint tracker instance
+	 */
+	setCheckpointTracker(checkpointTracker: CheckpointTracker | undefined): void {
+		this.state.checkpointTracker = checkpointTracker
+	}
+
+	/**
+	 * Updates the checkpoint tracker error message
+	 */
+	setcheckpointManagerErrorMessage(errorMessage: string | undefined): void {
+		this.state.checkpointManagerErrorMessage = errorMessage
+		// TODO - Future telemetry event capture here
+	}
+
+	/**
+	 * Updates the conversation history deleted range
+	 */
+	updateConversationHistoryDeletedRange(range: [number, number] | undefined): void {
+		this.state.conversationHistoryDeletedRange = range
+		// TODO - Future telemetry event capture here
+	}
+
+	// ============================================================================
+	// Internal utilities - Private helpers for checkpoint operations
+	// ============================================================================
+
+	/**
+	 * Gets the extension context with proper error handling
+	 */
+	private getContext(): vscode.ExtensionContext {
+		if (!this.services.context) {
+			throw new Error("Unable to access extension context")
+		}
+		return this.services.context
+	}
+
+	/**
+	 * Provides read-only access to current state for internal operations
+	 */
+	//private get currentState(): Readonly<CheckpointManagerInternalState> {
+	//	return Object.freeze({ ...this.state })
+	//}
+
+	/**
+	 * Provides public read-only access to current state
+	 */
+	public getCurrentState(): Readonly<CheckpointManagerInternalState> {
+		return Object.freeze({ ...this.state })
+	}
+
+	/**
+	 * Provides read-only access to dependencies for internal operations
+	 */
+	//private get deps(): Readonly<CheckpointManagerDependencies> {
+	//	return this.dependencies
+	//}
+}
+
+// ============================================================================
+// Factory function for clean instantiation
+// ============================================================================
+
+/**
+ * Creates a new TaskCheckpointManager instance
+ */
+export function createTaskCheckpointManager(
+	task: CheckpointManagerTask,
+	config: CheckpointManagerConfig,
+	services: CheckpointManagerServices,
+	callbacks: CheckpointManagerCallbacks,
+	initialState: CheckpointManagerInternalState,
+): TaskCheckpointManager {
+	return new TaskCheckpointManager(task, config, services, callbacks, initialState)
+}