Add completion change actions (Zoo-Code-Org#633)

* feat: add completion change actions

* fix: address completion action review feedback

* test(webview): mock console output and assert on checkpoint timeout error

- Mock `console.log` and `console.error` in `ClineProvider.flicker-free-cancel.spec.ts` to prevent test output pollution.
- Add assertion for `vscode.window.showErrorMessage` with `"errors.checkpoint_timeout"` in `webviewMessageHandler.checkpoint.spec.ts`.

* fix: address completion checkpoint review feedback

* feat(ChatRow): adding key to react component

---------

Co-authored-by: Elliott de Launay <edelauna@gmail.com>
# Conflicts:
#	webview-ui/src/components/chat/ChatView.tsx

Author: 2 people

README.md

@@ -87,7 +87,7 @@ for this exact support, so if you are having problems or if you have question, j
 - [简体中文](locales/zh-CN/README.md)
 - [繁體中文](locales/zh-TW/README.md)
 - ...
-    </details>
+  </details>
 
 ---
 

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: ClineMessage[] = [
+			{ 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 },
+		]
+
+		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

@@ -150,6 +150,7 @@ export const clineSays = [
 	"api_req_rate_limit_wait",
 	"api_req_deleted",
 	"text",
+	"task",
 	"image",
 	"reasoning",
 	"completion_result",
@@ -275,6 +276,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.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

@@ -520,6 +520,8 @@ export interface WebviewMessage {
 		| "openCustomModesSettings"
 		| "checkpointDiff"
 		| "checkpointRestore"
+		| "completionCheckpointDiff"
+		| "completionCheckpointRestore"
 		| "deleteMcpServer"
 		| "codebaseIndexEnabled"
 		| "telemetrySetting"

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

@@ -1,4 +1,4 @@
-import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"
+import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, vi } from "vitest"
 import * as vscode from "vscode"
 
 import { ClineProvider } from "../ClineProvider"
@@ -257,12 +257,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,109 @@ 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" })
+
+			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 () => {
+			const callOrder: string[] = []
+			mockProvider.cancelTask.mockImplementation(async () => callOrder.push("cancelTask"))
+			mockCline.checkpointRestore.mockImplementation(async () => callOrder.push("checkpointRestore"))
+
+			await webviewMessageHandler(mockProvider, { type: "completionCheckpointRestore" })
+
+			expect(mockProvider.cancelTask).toHaveBeenCalled()
+			expect(mockCline.checkpointRestore).toHaveBeenCalledWith({
+				ts: 4,
+				commitHash: "latest-prompt-checkpoint",
+				mode: "restore",
+			})
+			expect(callOrder).toEqual(["cancelTask", "checkpointRestore"])
+		})
+
+		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("resolves the latest completion checkpoint in the extension host", async () => {
+			await webviewMessageHandler(mockProvider, { type: "completionCheckpointDiff" })
+
+			expect(mockCline.checkpointDiff).toHaveBeenCalledWith({
+				ts: 4,
+				commitHash: "latest-prompt-checkpoint",
+				mode: "to-current",
+			})
+		})
+
+		it("does not restore when task re-initialization times out", async () => {
+			;(pWaitFor as any).mockRejectedValueOnce(new Error("timed out"))
+
+			await webviewMessageHandler(mockProvider, { type: "completionCheckpointRestore" })
+
+			expect(mockProvider.cancelTask).toHaveBeenCalled()
+			expect(mockCline.checkpointRestore).not.toHaveBeenCalled()
+			const vscode = await import("vscode")
+			expect(vscode.window.showErrorMessage).toHaveBeenCalledWith("errors.checkpoint_timeout")
+		})
+
+		it("shows an error when completion checkpoint restore fails", async () => {
+			const restoreError = new Error("restore failed")
+			const consoleErrorSpy = vi.spyOn(console, "error").mockImplementation(() => {})
+			mockCline.checkpointRestore.mockRejectedValueOnce(restoreError)
+
+			await webviewMessageHandler(mockProvider, { type: "completionCheckpointRestore" })
+
+			const vscode = await import("vscode")
+			expect(consoleErrorSpy).toHaveBeenCalledWith(
+				"[completionCheckpointRestore] checkpointRestore failed:",
+				restoreError,
+			)
+			expect(vscode.window.showErrorMessage).toHaveBeenCalledWith("errors.checkpoint_failed")
+			consoleErrorSpy.mockRestore()
+		})
+
+		it("does not restore when task identity changes during cancellation", async () => {
+			mockProvider.getCurrentTask.mockReturnValueOnce(mockCline).mockReturnValue({
+				...mockCline,
+				taskId: "different-task-id",
+			})
+
+			await webviewMessageHandler(mockProvider, { type: "completionCheckpointRestore" })
+
+			expect(mockProvider.cancelTask).toHaveBeenCalled()
+			expect(mockCline.checkpointRestore).not.toHaveBeenCalled()
+			const vscode = await import("vscode")
+			expect(vscode.window.showErrorMessage).toHaveBeenCalledWith("errors.checkpoint_failed")
+		})
+	})
 })