Add completion change actions

packages/types/src/__tests__/message.test.ts

@@ -1,6 +1,14 @@
 // pnpm --filter @roo-code/types test src/__tests__/message.test.ts
 
-import { clineAsks, isIdleAsk, isInteractiveAsk, isResumableAsk, isNonBlockingAsk } from "../message.js"
+import {
+	clineAsks,
+	getCompletionCheckpoint,
+	isIdleAsk,
+	isInteractiveAsk,
+	isResumableAsk,
+	isNonBlockingAsk,
+	type ClineMessage,
+} from "../message.js"
 
 describe("ask messages", () => {
 	test("all ask messages are classified", () => {
@@ -12,3 +20,46 @@ describe("ask messages", () => {
 		}
 	})
 })
+
+describe("getCompletionCheckpoint", () => {
+	it("returns the first checkpoint after the latest user prompt before completion", () => {
+		const messages: ClineMessage[] = [
+			{ type: "say", say: "text", ts: 1, text: "Initial task" },
+			{ type: "say", say: "checkpoint_saved", ts: 2, text: "initial-checkpoint" },
+			{ type: "say", say: "completion_result", ts: 3, text: "First completion" },
+			{ type: "say", say: "user_feedback", ts: 4, text: "Change it" },
+			{ type: "say", say: "checkpoint_saved", ts: 5, text: "latest-prompt-checkpoint" },
+			{ type: "say", say: "checkpoint_saved", ts: 6, text: "later-edit-checkpoint" },
+			{ type: "ask", ask: "completion_result", ts: 7, text: "", partial: false },
+		]
+
+		expect(getCompletionCheckpoint(messages)).toEqual({
+			ts: 5,
+			commitHash: "latest-prompt-checkpoint",
+		})
+	})
+
+	it("returns the first checkpoint after an initial task row before completion", () => {
+		const messages = [
+			{ type: "say", say: "task", ts: 1, text: "Initial task" },
+			{ type: "say", say: "checkpoint_saved", ts: 2, text: "checkpoint-after-initial-task" },
+			{ type: "ask", ask: "completion_result", ts: 3, text: "Task complete", partial: false },
+		] as ClineMessage[]
+
+		expect(getCompletionCheckpoint(messages)).toEqual({
+			ts: 2,
+			commitHash: "checkpoint-after-initial-task",
+		})
+	})
+
+	it("returns undefined when completion has no checkpoint after the latest user prompt", () => {
+		const messages: ClineMessage[] = [
+			{ type: "say", say: "text", ts: 1, text: "Initial task" },
+			{ type: "say", say: "checkpoint_saved", ts: 2, text: "initial-checkpoint" },
+			{ type: "say", say: "user_feedback", ts: 3, text: "Change it" },
+			{ type: "ask", ask: "completion_result", ts: 4, text: "", partial: false },
+		]
+
+		expect(getCompletionCheckpoint(messages)).toBeUndefined()
+	})
+})

packages/types/src/message.ts

@@ -275,6 +275,76 @@ export const clineMessageSchema = z.object({
 
 export type ClineMessage = z.infer<typeof clineMessageSchema>
 
+export interface CompletionCheckpoint {
+	ts: number
+	commitHash: string
+}
+
+const isInitialTaskMessage = (message: ClineMessage | undefined): boolean => {
+	return message?.type === "say" && (message.say === "text" || (message as { say?: string }).say === "task")
+}
+
+const isUserFeedbackMessage = (message: ClineMessage): boolean => {
+	return message.type === "say" && message.say === "user_feedback"
+}
+
+const isCompletionMessage = (message: ClineMessage): boolean => {
+	return (
+		(message.type === "ask" && message.ask === "completion_result") ||
+		(message.type === "say" && message.say === "completion_result")
+	)
+}
+
+const isCheckpointMessage = (message: ClineMessage): boolean => {
+	return message.type === "say" && message.say === "checkpoint_saved" && typeof message.text === "string"
+}
+
+function findLastIndexBefore(
+	messages: ClineMessage[],
+	beforeIndex: number,
+	predicate: (message: ClineMessage) => boolean,
+): number {
+	for (let i = beforeIndex - 1; i >= 0; i--) {
+		const message = messages[i]
+
+		if (message && predicate(message)) {
+			return i
+		}
+	}
+
+	return -1
+}
+
+/**
+ * Finds the checkpoint that should anchor completion-result actions.
+ *
+ * The baseline is the first checkpoint created after the latest user prompt in
+ * the turn that produced the completion. Restoring to that checkpoint reverts
+ * changes made for the latest prompt, and diffing from it shows the same scoped
+ * changes.
+ */
+export function getCompletionCheckpoint(messages: ClineMessage[]): CompletionCheckpoint | undefined {
+	const completionIndex = findLastIndexBefore(messages, messages.length, isCompletionMessage)
+	const searchEnd = completionIndex === -1 ? messages.length : completionIndex
+	const latestUserFeedbackIndex = findLastIndexBefore(messages, searchEnd, isUserFeedbackMessage)
+	const latestUserPromptIndex =
+		latestUserFeedbackIndex !== -1 ? latestUserFeedbackIndex : isInitialTaskMessage(messages[0]) ? 0 : -1
+
+	if (latestUserPromptIndex === -1) {
+		return undefined
+	}
+
+	for (let i = latestUserPromptIndex + 1; i < searchEnd; i++) {
+		const message = messages[i]
+
+		if (message && isCheckpointMessage(message)) {
+			return { ts: message.ts, commitHash: message.text! }
+		}
+	}
+
+	return undefined
+}
+
 /**
  * TokenUsage
  */

packages/types/src/vscode-extension-host.ts

@@ -507,6 +507,8 @@ export interface WebviewMessage {
 		| "openCustomModesSettings"
 		| "checkpointDiff"
 		| "checkpointRestore"
+		| "completionCheckpointDiff"
+		| "completionCheckpointRestore"
 		| "deleteMcpServer"
 		| "codebaseIndexEnabled"
 		| "telemetrySetting"
@@ -646,6 +648,7 @@ export interface WebviewMessage {
 	ids?: string[]
 	terminalOperation?: "continue" | "abort"
 	messageTs?: number
+	checkpointTs?: number
 	restoreCheckpoint?: boolean
 	historyPreviewCollapsed?: boolean
 	filters?: { type?: string; search?: string; tags?: string[] }

src/core/webview/__tests__/ClineProvider.flicker-free-cancel.spec.ts

@@ -1,4 +1,4 @@
-import { beforeEach, describe, expect, it, vi } from "vitest"
+import { afterAll, beforeAll, beforeEach, describe, expect, it, vi } from "vitest"
 import * as vscode from "vscode"
 
 import { ClineProvider } from "../ClineProvider"
@@ -93,12 +93,24 @@ describe("ClineProvider flicker-free cancel", () => {
 	let mockOutputChannel: any
 	let mockTask1: any
 	let mockTask2: any
+	let consoleLogSpy: ReturnType<typeof vi.spyOn>
+	let consoleErrorSpy: ReturnType<typeof vi.spyOn>
 
 	const mockApiConfig: ProviderSettings = {
 		apiProvider: "anthropic",
 		apiKey: "test-key",
 	} as ProviderSettings
 
+	beforeAll(() => {
+		consoleLogSpy = vi.spyOn(console, "log").mockImplementation(() => {})
+		consoleErrorSpy = vi.spyOn(console, "error").mockImplementation(() => {})
+	})
+
+	afterAll(() => {
+		consoleLogSpy.mockRestore()
+		consoleErrorSpy.mockRestore()
+	})
+
 	beforeEach(() => {
 		vi.clearAllMocks()
 

src/core/webview/__tests__/webviewMessageHandler.checkpoint.spec.ts

@@ -1,4 +1,5 @@
 import { describe, it, expect, vi, beforeEach } from "vitest"
+import pWaitFor from "p-wait-for"
 import { webviewMessageHandler } from "../webviewMessageHandler"
 import { saveTaskMessages } from "../../task-persistence"
 import { handleCheckpointRestoreOperation } from "../checkpointRestoreHandler"
@@ -7,6 +8,13 @@ import { MessageManager } from "../../message-manager"
 // Mock dependencies
 vi.mock("../../task-persistence")
 vi.mock("../checkpointRestoreHandler")
+vi.mock("p-wait-for", () => ({
+	default: vi.fn(async (condition: () => boolean) => {
+		if (!condition()) {
+			throw new Error("condition not met")
+		}
+	}),
+}))
 vi.mock("vscode", () => ({
 	window: {
 		showErrorMessage: vi.fn(),
@@ -26,6 +34,7 @@ describe("webviewMessageHandler - checkpoint operations", () => {
 		// Setup mock Cline instance
 		mockCline = {
 			taskId: "test-task-123",
+			isInitialized: true,
 			clineMessages: [
 				{ ts: 1, type: "user", say: "user", text: "First message" },
 				{ ts: 2, type: "assistant", say: "checkpoint_saved", text: "abc123" },
@@ -37,6 +46,7 @@ describe("webviewMessageHandler - checkpoint operations", () => {
 				{ ts: 3, role: "user", content: [{ type: "text", text: "Message to delete" }] },
 				{ ts: 4, role: "assistant", content: [{ type: "text", text: "After message" }] },
 			],
+			checkpointDiff: vi.fn(),
 			checkpointRestore: vi.fn(),
 			overwriteClineMessages: vi.fn(),
 			overwriteApiConversationHistory: vi.fn(),
@@ -52,6 +62,7 @@ describe("webviewMessageHandler - checkpoint operations", () => {
 			})),
 			createTaskWithHistoryItem: vi.fn(),
 			setPendingEditOperation: vi.fn(),
+			cancelTask: vi.fn(),
 			contextProxy: {
 				globalStorageUri: { fsPath: "/test/storage" },
 			},
@@ -134,4 +145,73 @@ describe("webviewMessageHandler - checkpoint operations", () => {
 			})
 		})
 	})
+
+	describe("completion checkpoint actions", () => {
+		beforeEach(() => {
+			mockCline.clineMessages = [
+				{ ts: 1, type: "say", say: "text", text: "Initial task" },
+				{ ts: 2, type: "say", say: "checkpoint_saved", text: "initial-checkpoint" },
+				{ ts: 3, type: "say", say: "user_feedback", text: "Latest prompt" },
+				{ ts: 4, type: "say", say: "checkpoint_saved", text: "latest-prompt-checkpoint" },
+				{ ts: 5, type: "say", say: "completion_result", text: "Task complete" },
+				{ ts: 6, type: "ask", ask: "completion_result", text: "", partial: false },
+			]
+		})
+
+		it("diffs changes from the checkpoint created after the latest prompt", async () => {
+			await webviewMessageHandler(mockProvider, { type: "completionCheckpointDiff", checkpointTs: 4 })
+
+			expect(mockCline.checkpointDiff).toHaveBeenCalledWith({
+				ts: 4,
+				commitHash: "latest-prompt-checkpoint",
+				mode: "to-current",
+			})
+		})
+
+		it("restores files and task state to the checkpoint created after the latest prompt", async () => {
+			await webviewMessageHandler(mockProvider, { type: "completionCheckpointRestore", checkpointTs: 4 })
+
+			expect(mockProvider.cancelTask).toHaveBeenCalled()
+			expect(mockCline.checkpointRestore).toHaveBeenCalledWith({
+				ts: 4,
+				commitHash: "latest-prompt-checkpoint",
+				mode: "restore",
+			})
+		})
+
+		it("does not diff or restore when no latest-prompt checkpoint exists", async () => {
+			mockCline.clineMessages = [
+				{ ts: 1, type: "say", say: "text", text: "Initial task" },
+				{ ts: 2, type: "say", say: "user_feedback", text: "Latest prompt" },
+				{ ts: 3, type: "ask", ask: "completion_result", text: "", partial: false },
+			]
+
+			await webviewMessageHandler(mockProvider, { type: "completionCheckpointDiff" })
+			await webviewMessageHandler(mockProvider, { type: "completionCheckpointRestore" })
+
+			expect(mockCline.checkpointDiff).not.toHaveBeenCalled()
+			expect(mockCline.checkpointRestore).not.toHaveBeenCalled()
+			expect(mockProvider.cancelTask).not.toHaveBeenCalled()
+		})
+
+		it("does not re-derive a different checkpoint when an explicit checkpoint timestamp is missing", async () => {
+			await webviewMessageHandler(mockProvider, { type: "completionCheckpointDiff", checkpointTs: 999 })
+			await webviewMessageHandler(mockProvider, { type: "completionCheckpointRestore", checkpointTs: 999 })
+
+			expect(mockCline.checkpointDiff).not.toHaveBeenCalled()
+			expect(mockCline.checkpointRestore).not.toHaveBeenCalled()
+			expect(mockProvider.cancelTask).not.toHaveBeenCalled()
+		})
+
+		it("does not restore when task re-initialization times out", async () => {
+			;(pWaitFor as any).mockRejectedValueOnce(new Error("timed out"))
+
+			await webviewMessageHandler(mockProvider, { type: "completionCheckpointRestore", checkpointTs: 4 })
+
+			expect(mockProvider.cancelTask).toHaveBeenCalled()
+			expect(mockCline.checkpointRestore).not.toHaveBeenCalled()
+			const vscode = await import("vscode")
+			expect(vscode.window.showErrorMessage).toHaveBeenCalledWith("errors.checkpoint_timeout")
+		})
+	})
 })

src/core/webview/webviewMessageHandler.ts

@@ -21,6 +21,7 @@ import {
 	ExperimentId,
 	checkoutDiffPayloadSchema,
 	checkoutRestorePayloadSchema,
+	getCompletionCheckpoint,
 } from "@roo-code/types"
 import { customToolRegistry } from "@roo-code/core"
 import { CloudService } from "@roo-code/cloud"
@@ -112,6 +113,25 @@ export const webviewMessageHandler = async (
 		vscode.window.showInformationMessage(getRouterUnavailableSignInMessage())
 	}
 
+	const getCompletionCheckpointForMessage = (
+		currentCline: { clineMessages: ClineMessage[] },
+		checkpointTs?: number,
+	) => {
+		if (checkpointTs !== undefined) {
+			const checkpoint = currentCline.clineMessages.find(
+				(message) =>
+					message.ts === checkpointTs &&
+					message.type === "say" &&
+					message.say === "checkpoint_saved" &&
+					typeof message.text === "string",
+			)
+
+			return checkpoint ? { ts: checkpoint.ts, commitHash: checkpoint.text! } : undefined
+		}
+
+		return getCompletionCheckpoint(currentCline.clineMessages)
+	}
+
 	const getCurrentMode = async (): Promise<string> => {
 		const currentTask = provider.getCurrentTask()
 
@@ -1288,6 +1308,51 @@ export const webviewMessageHandler = async (
 
 			break
 		}
+		case "completionCheckpointDiff": {
+			const currentCline = provider.getCurrentTask()
+			const checkpoint = currentCline
+				? getCompletionCheckpointForMessage(currentCline, message.checkpointTs)
+				: undefined
+
+			if (currentCline && checkpoint) {
+				await currentCline.checkpointDiff({
+					ts: checkpoint.ts,
+					commitHash: checkpoint.commitHash,
+					mode: "to-current",
+				})
+			}
+
+			break
+		}
+		case "completionCheckpointRestore": {
+			const currentCline = provider.getCurrentTask()
+			const checkpoint = currentCline
+				? getCompletionCheckpointForMessage(currentCline, message.checkpointTs)
+				: undefined
+
+			if (checkpoint) {
+				await provider.cancelTask()
+
+				try {
+					await pWaitFor(() => provider.getCurrentTask()?.isInitialized === true, { timeout: 3_000 })
+				} catch (error) {
+					vscode.window.showErrorMessage(t("common:errors.checkpoint_timeout"))
+					return
+				}
+
+				try {
+					await provider.getCurrentTask()?.checkpointRestore({
+						ts: checkpoint.ts,
+						commitHash: checkpoint.commitHash,
+						mode: "restore",
+					})
+				} catch (error) {
+					vscode.window.showErrorMessage(t("common:errors.checkpoint_failed"))
+				}
+			}
+
+			break
+		}
 		case "cancelTask":
 			await provider.cancelTask()
 			break