Merge branch 'main' into pm/unbound-fetch-models

Author: vigneshsubbiah16

.gitignore

@@ -22,3 +22,4 @@ docs/_site/
 
 # Dotenv
 .env.integration
+.eslintrc.local.json

package.json

@@ -272,6 +272,7 @@
 		"compile:integration": "tsc -p tsconfig.integration.json",
 		"install:all": "npm install && cd webview-ui && npm install",
 		"lint": "eslint src --ext ts && npm run lint --prefix webview-ui",
+		"lint-fix": "eslint src --ext ts --fix && npm run lint-fix --prefix webview-ui",
 		"package": "npm run build:webview && npm run check-types && npm run lint && node esbuild.js --production",
 		"pretest": "npm run compile && npm run compile:integration",
 		"dev": "cd webview-ui && npm run dev",
@@ -283,7 +284,7 @@
 		"publish": "npm run build && changeset publish && npm install --package-lock-only",
 		"version-packages": "changeset version && npm install --package-lock-only",
 		"vscode:prepublish": "npm run package",
-		"vsix": "mkdir -p bin && npx vsce package --out bin",
+		"vsix": "rimraf bin && mkdirp bin && npx vsce package --out bin",
 		"watch": "npm-run-all -p watch:*",
 		"watch:esbuild": "node esbuild.js --watch",
 		"watch:tsc": "tsc --noEmit --watch --project tsconfig.json",
@@ -350,6 +351,8 @@
 		"@vscode/test-cli": "^0.0.9",
 		"@vscode/test-electron": "^2.4.0",
 		"esbuild": "^0.24.0",
+		"mkdirp": "^3.0.1",
+		"rimraf": "^6.0.1",
 		"eslint": "^8.57.0",
 		"husky": "^9.1.7",
 		"jest": "^29.7.0",

src/core/Cline.ts

@@ -11,7 +11,8 @@ import { serializeError } from "serialize-error"
 import * as vscode from "vscode"
 import { ApiHandler, SingleCompletionHandler, buildApiHandler } from "../api"
 import { ApiStream } from "../api/transform/stream"
-import { DiffViewProvider } from "../integrations/editor/DiffViewProvider"
+import { DIFF_VIEW_URI_SCHEME, DiffViewProvider } from "../integrations/editor/DiffViewProvider"
+import { CheckpointService } from "../services/checkpoints/CheckpointService"
 import { findToolName, formatContentBlockToMarkdown } from "../integrations/misc/export-markdown"
 import {
 	extractTextFromFile,
@@ -93,12 +94,19 @@ export class Cline {
 	private consecutiveMistakeCountForApplyDiff: Map<string, number> = new Map()
 	private providerRef: WeakRef<ClineProvider>
 	private abort: boolean = false
-	didFinishAborting = false
+	didFinishAbortingStream = false
 	abandoned = false
 	private diffViewProvider: DiffViewProvider
 	private lastApiRequestTime?: number
+	isInitialized = false
+
+	// checkpoints
+	checkpointsEnabled: boolean = false
+	private checkpointService?: CheckpointService
 
 	// streaming
+	isWaitingForFirstChunk = false
+	isStreaming = false
 	private currentStreamingContentIndex = 0
 	private assistantMessageContent: AssistantMessageContent[] = []
 	private presentAssistantMessageLocked = false
@@ -114,6 +122,7 @@ export class Cline {
 		apiConfiguration: ApiConfiguration,
 		customInstructions?: string,
 		enableDiff?: boolean,
+		enableCheckpoints?: boolean,
 		fuzzyMatchThreshold?: number,
 		task?: string | undefined,
 		images?: string[] | undefined,
@@ -134,6 +143,7 @@ export class Cline {
 		this.fuzzyMatchThreshold = fuzzyMatchThreshold ?? 1.0
 		this.providerRef = new WeakRef(provider)
 		this.diffViewProvider = new DiffViewProvider(cwd)
+		this.checkpointsEnabled = enableCheckpoints ?? false
 
 		if (historyItem) {
 			this.taskId = historyItem.id
@@ -438,6 +448,7 @@ export class Cline {
 		await this.providerRef.deref()?.postStateToWebview()
 
 		await this.say("text", task, images)
+		this.isInitialized = true
 
 		let imageBlocks: Anthropic.ImageBlockParam[] = formatResponse.imageBlocks(images)
 		await this.initiateTaskLoop([
@@ -477,12 +488,13 @@ export class Cline {
 		await this.overwriteClineMessages(modifiedClineMessages)
 		this.clineMessages = await this.getSavedClineMessages()
 
-		// need to 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: Anthropic.Messages.MessageParam[] =
-			await this.getSavedApiConversationHistory()
-
-		// Now present the cline messages to the user and ask if they want to resume
+		// Now present the cline messages to the user and ask if they want to
+		// resume (NOTE: we ran into a bug before where the
+		// apiConversationHistory wouldn't be initialized when opening a old
+		// task, and it was because we were waiting for resume).
+		// This is important in case the user deletes messages without resuming
+		// the task first.
+		this.apiConversationHistory = await this.getSavedApiConversationHistory()
 
 		const lastClineMessage = this.clineMessages
 			.slice()
@@ -506,6 +518,8 @@ export class Cline {
 			askType = "resume_task"
 		}
 
+		this.isInitialized = true
+
 		const { response, text, images } = await this.ask(askType) // calls poststatetowebview
 		let responseText: string | undefined
 		let responseImages: string[] | undefined
@@ -515,6 +529,11 @@ export class Cline {
 			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: Anthropic.Messages.MessageParam[] =
+			await this.getSavedApiConversationHistory()
+
 		// v2.0 xml tags refactor caveat: since we don't use tools anymore, we need to replace all tool use blocks with a text block since the API disallows conversations with tool uses and no tool schema
 		const conversationWithoutToolBlocks = existingApiConversationHistory.map((message) => {
 			if (Array.isArray(message.content)) {
@@ -706,11 +725,14 @@ export class Cline {
 		}
 	}
 
-	abortTask() {
-		this.abort = true // will stop any autonomously running promises
+	async abortTask() {
+		this.abort = true // Will stop any autonomously running promises.
 		this.terminalManager.disposeAll()
 		this.urlContentFetcher.closeBrowser()
 		this.browserSession.closeBrowser()
+		// Need to await for when we want to make sure directories/files are
+		// reverted before re-starting the task from a checkpoint.
+		await this.diffViewProvider.revertChanges()
 	}
 
 	// Tools
@@ -927,8 +949,10 @@ export class Cline {
 
 		try {
 			// awaiting first chunk to see if it will throw an error
+			this.isWaitingForFirstChunk = true
 			const firstChunk = await iterator.next()
 			yield firstChunk.value
+			this.isWaitingForFirstChunk = false
 		} catch (error) {
 			// note that this api_req_failed ask is unique in that we only present this option if the api hasn't streamed any content yet (ie it fails on the first chunk due), as it would allow them to hit a retry button. However if the api failed mid-stream, it could be in any arbitrary state where some tools may have executed, so that error is handled differently and requires cancelling the task entirely.
 			if (alwaysApproveResubmit) {
@@ -1003,6 +1027,9 @@ export class Cline {
 		}
 
 		const block = cloneDeep(this.assistantMessageContent[this.currentStreamingContentIndex]) // need to create copy bc while stream is updating the array, it could be updating the reference block properties too
+
+		let isCheckpointPossible = false
+
 		switch (block.type) {
 			case "text": {
 				if (this.didRejectTool || this.didAlreadyUseTool) {
@@ -1134,6 +1161,10 @@ export class Cline {
 					}
 					// once a tool result has been collected, ignore all other tool uses since we should only ever present one tool result per message
 					this.didAlreadyUseTool = true
+
+					// Flag a checkpoint as possible since we've used a tool
+					// which may have changed the file system.
+					isCheckpointPossible = true
 				}
 
 				const askApproval = async (type: ClineAsk, partialMessage?: string) => {
@@ -2655,6 +2686,10 @@ export class Cline {
 				break
 		}
 
+		if (isCheckpointPossible) {
+			await this.checkpointSave()
+		}
+
 		/*
 		Seeing out of bounds is fine, it means that the next too call is being built up and ready to add to assistantMessageContent to present.
 		When you see the UI inactive during this, it means that a tool is breaking without presenting any UI. For example the write_to_file tool was breaking when relpath was undefined, and for invalid relpath it never presented UI.
@@ -2811,7 +2846,7 @@ export class Cline {
 				await this.saveClineMessages()
 
 				// signals to provider that it can retrieve the saved messages from disk, as abortTask can not be awaited on in nature
-				this.didFinishAborting = true
+				this.didFinishAbortingStream = true
 			}
 
 			// reset streaming state
@@ -3197,6 +3232,178 @@ export class Cline {
 
 		return `<environment_details>\n${details.trim()}\n</environment_details>`
 	}
+
+	// Checkpoints
+
+	private async getCheckpointService() {
+		if (!this.checkpointService) {
+			this.checkpointService = await CheckpointService.create({
+				taskId: this.taskId,
+				baseDir: vscode.workspace.workspaceFolders?.map((folder) => folder.uri.fsPath).at(0) ?? "",
+			})
+		}
+
+		return this.checkpointService
+	}
+
+	public async checkpointDiff({
+		ts,
+		commitHash,
+		mode,
+	}: {
+		ts: number
+		commitHash: string
+		mode: "full" | "checkpoint"
+	}) {
+		if (!this.checkpointsEnabled) {
+			return
+		}
+
+		let previousCommitHash = undefined
+
+		if (mode === "checkpoint") {
+			const previousCheckpoint = this.clineMessages
+				.filter(({ say }) => say === "checkpoint_saved")
+				.sort((a, b) => b.ts - a.ts)
+				.find((message) => message.ts < ts)
+
+			previousCommitHash = previousCheckpoint?.text
+		}
+
+		try {
+			const service = await this.getCheckpointService()
+			const changes = await service.getDiff({ from: previousCommitHash, to: commitHash })
+
+			if (!changes?.length) {
+				vscode.window.showInformationMessage("No changes found.")
+				return
+			}
+
+			await vscode.commands.executeCommand(
+				"vscode.changes",
+				mode === "full" ? "Changes since task started" : "Changes since previous checkpoint",
+				changes.map((change) => [
+					vscode.Uri.file(change.paths.absolute),
+					vscode.Uri.parse(`${DIFF_VIEW_URI_SCHEME}:${change.paths.relative}`).with({
+						query: Buffer.from(change.content.before ?? "").toString("base64"),
+					}),
+					vscode.Uri.parse(`${DIFF_VIEW_URI_SCHEME}:${change.paths.relative}`).with({
+						query: Buffer.from(change.content.after ?? "").toString("base64"),
+					}),
+				]),
+			)
+		} catch (err) {
+			this.providerRef
+				.deref()
+				?.log(
+					`[checkpointDiff] Encountered unexpected error: $${err instanceof Error ? err.message : String(err)}`,
+				)
+
+			this.checkpointsEnabled = false
+		}
+	}
+
+	public async checkpointSave() {
+		if (!this.checkpointsEnabled) {
+			return
+		}
+
+		try {
+			const service = await this.getCheckpointService()
+			const commit = await service.saveCheckpoint(`Task: ${this.taskId}, Time: ${Date.now()}`)
+
+			if (commit?.commit) {
+				await this.providerRef
+					.deref()
+					?.postMessageToWebview({ type: "currentCheckpointUpdated", text: service.currentCheckpoint })
+
+				await this.say("checkpoint_saved", commit.commit)
+			}
+		} catch (err) {
+			this.providerRef
+				.deref()
+				?.log(
+					`[checkpointSave] Encountered unexpected error: $${err instanceof Error ? err.message : String(err)}`,
+				)
+
+			this.checkpointsEnabled = false
+		}
+	}
+
+	public async checkpointRestore({
+		ts,
+		commitHash,
+		mode,
+	}: {
+		ts: number
+		commitHash: string
+		mode: "preview" | "restore"
+	}) {
+		if (!this.checkpointsEnabled) {
+			return
+		}
+
+		const index = this.clineMessages.findIndex((m) => m.ts === ts)
+
+		if (index === -1) {
+			return
+		}
+
+		try {
+			const service = await this.getCheckpointService()
+			await service.restoreCheckpoint(commitHash)
+
+			await this.providerRef
+				.deref()
+				?.postMessageToWebview({ type: "currentCheckpointUpdated", text: service.currentCheckpoint })
+
+			if (mode === "restore") {
+				await this.overwriteApiConversationHistory(
+					this.apiConversationHistory.filter((m) => !m.ts || m.ts < ts),
+				)
+
+				const deletedMessages = this.clineMessages.slice(index + 1)
+
+				const { totalTokensIn, totalTokensOut, totalCacheWrites, totalCacheReads, totalCost } = getApiMetrics(
+					combineApiRequests(combineCommandSequences(deletedMessages)),
+				)
+
+				await this.overwriteClineMessages(this.clineMessages.slice(0, index + 1))
+
+				// TODO: Verify that this is working as expected.
+				await this.say(
+					"api_req_deleted",
+					JSON.stringify({
+						tokensIn: totalTokensIn,
+						tokensOut: totalTokensOut,
+						cacheWrites: totalCacheWrites,
+						cacheReads: totalCacheReads,
+						cost: totalCost,
+					} satisfies ClineApiReqInfo),
+				)
+			}
+
+			// The task is already cancelled by the provider beforehand, but we
+			// need to re-init to get the updated messages.
+			//
+			// This was take from Cline's implementation of the checkpoints
+			// feature. The cline instance will hang if we don't cancel twice,
+			// so this is currently necessary, but it seems like a complicated
+			// and hacky solution to a problem that I don't fully understand.
+			// I'd like to revisit this in the future and try to improve the
+			// task flow and the communication between the webview and the
+			// Cline instance.
+			this.providerRef.deref()?.cancelTask()
+		} catch (err) {
+			this.providerRef
+				.deref()
+				?.log(
+					`[restoreCheckpoint] Encountered unexpected error: $${err instanceof Error ? err.message : String(err)}`,
+				)
+
+			this.checkpointsEnabled = false
+		}
+	}
 }
 
 function escapeRegExp(string: string): string {