Add completion change actions

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"
+}
+
+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

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

src/core/webview/__tests__/completionCheckpoint.spec.ts

@@ -0,0 +1,31 @@
+import { getCompletionCheckpoint, type ClineMessage } from "@roo-code/types"
+
+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 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()
+	})
+})

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

@@ -26,6 +26,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 +38,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 +54,7 @@ describe("webviewMessageHandler - checkpoint operations", () => {
 			})),
 			createTaskWithHistoryItem: vi.fn(),
 			setPendingEditOperation: vi.fn(),
+			cancelTask: vi.fn(),
 			contextProxy: {
 				globalStorageUri: { fsPath: "/test/storage" },
 			},
@@ -134,4 +137,53 @@ 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 () => {
+			await webviewMessageHandler(mockProvider, { type: "completionCheckpointRestore" })
+
+			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()
+		})
+	})
 })

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"
@@ -1288,6 +1289,46 @@ export const webviewMessageHandler = async (
 
 			break
 		}
+		case "completionCheckpointDiff": {
+			const currentCline = provider.getCurrentTask()
+			const checkpoint = currentCline ? getCompletionCheckpoint(currentCline.clineMessages) : 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 ? getCompletionCheckpoint(currentCline.clineMessages) : 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"))
+				}
+
+				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

webview-ui/src/components/chat/ChatRow.tsx

@@ -11,6 +11,7 @@ import type {
 	ClineApiReqInfo,
 	ClineAskUseMcpServer,
 	ClineSayTool,
+	CompletionCheckpoint,
 } from "@roo-code/types"
 
 import { Mode } from "@roo/modes"
@@ -75,6 +76,7 @@ import {
 import { cn } from "@/lib/utils"
 import { PathTooltip } from "../ui/PathTooltip"
 import { OpenMarkdownPreviewButton } from "./OpenMarkdownPreviewButton"
+import { SeeNewChangesButtons } from "./SeeNewChangesButtons"
 
 // Helper function to get previous todos before a specific message
 function getPreviousTodos(messages: ClineMessage[], currentMessageTs: number): any[] {
@@ -124,6 +126,7 @@ interface ChatRowProps {
 	isFollowUpAutoApprovalPaused?: boolean
 	editable?: boolean
 	hasCheckpoint?: boolean
+	completionCheckpoint?: CompletionCheckpoint
 	onJumpToPreviousCheckpoint?: () => void
 }
 
@@ -178,12 +181,21 @@ export const ChatRowContent = ({
 	onBatchFileResponse,
 	isFollowUpAnswered,
 	isFollowUpAutoApprovalPaused,
+	completionCheckpoint,
 	onJumpToPreviousCheckpoint,
 }: ChatRowContentProps) => {
 	const { t, i18n } = useTranslation()
 
-	const { mcpServers, alwaysAllowMcp, currentCheckpoint, mode, apiConfiguration, clineMessages, currentTaskItem } =
-		useExtensionState()
+	const {
+		mcpServers,
+		alwaysAllowMcp,
+		currentCheckpoint,
+		mode,
+		apiConfiguration,
+		clineMessages,
+		currentTaskItem,
+		enableCheckpoints,
+	} = useExtensionState()
 	const { info: model } = useSelectedModel(apiConfiguration)
 	const [isEditing, setIsEditing] = useState(false)
 	const [editedContent, setEditedContent] = useState("")
@@ -1331,6 +1343,9 @@ export const ChatRowContent = ({
 							</div>
 							<div className="border-l border-green-600/30 ml-2 pl-4 pb-1">
 								<Markdown markdown={message.text} />
+								{!message.partial && enableCheckpoints !== false && completionCheckpoint ? (
+									<SeeNewChangesButtons checkpoint={completionCheckpoint} />
+								) : null}
 							</div>
 						</div>
 					)

webview-ui/src/components/chat/ChatView.tsx

@@ -11,7 +11,7 @@ import { getCostBreakdownIfNeeded } from "@src/utils/costFormatting"
 import { batchConsecutive } from "@src/utils/batchConsecutive"
 
 import type { ClineAsk, ClineSayTool, ClineMessage, ExtensionMessage, AudioType, SuggestionItem } from "@roo-code/types"
-import { getSuggestionMode, isRetiredProvider } from "@roo-code/types"
+import { getCompletionCheckpoint, getSuggestionMode, isRetiredProvider } from "@roo-code/types"
 
 import { findLast } from "@roo/array"
 import { combineApiRequests } from "@roo/combineApiRequests"
@@ -129,6 +129,18 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
 	}, [messages, currentTaskTodos])
 
 	const modifiedMessages = useMemo(() => combineApiRequests(combineCommandSequences(messages.slice(1))), [messages])
+	const completionCheckpoint = useMemo(() => getCompletionCheckpoint(messages), [messages])
+	const completionResultTs = useMemo(() => {
+		for (let i = messages.length - 1; i >= 0; i--) {
+			const message = messages[i]
+
+			if (message?.type === "say" && message.say === "completion_result") {
+				return message.ts
+			}
+		}
+
+		return undefined
+	}, [messages])
 
 	// Has to be after api_req_finished are all reduced into api_req_started messages.
 	const apiMetrics = useMemo(() => getApiMetrics(modifiedMessages), [modifiedMessages])
@@ -374,6 +386,7 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
 							break
 						case "completion_result":
 							// Extension waiting for feedback, but we can just present a new task button.
+							// Kilo-style change inspection/restoration buttons are rendered inline on the completion row.
 							// Only play celebration sound if there are no queued messages.
 							if (!isPartial && messageQueue.length === 0) {
 								playSound("celebration")
@@ -788,6 +801,8 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
 					}
 					break
 				case "completion_result":
+					startNewTask()
+					break
 				case "resume_completed_task":
 					// Waiting for feedback, but we can just present a new task button
 					startNewTask()
@@ -1476,6 +1491,7 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
 						})()
 					}
 					hasCheckpoint={hasCheckpoint}
+					completionCheckpoint={messageOrGroup.ts === completionResultTs ? completionCheckpoint : undefined}
 					onJumpToPreviousCheckpoint={handleScrollToLatestCheckpoint}
 				/>
 			)
@@ -1485,6 +1501,8 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
 			toggleRowExpansion,
 			modifiedMessages,
 			groupedMessages.length,
+			completionCheckpoint,
+			completionResultTs,
 			handleRowHeightChange,
 			isStreaming,
 			handleSuggestionClickInRow,