diff --git a/packages/types/src/vscode-extension-host.ts b/packages/types/src/vscode-extension-host.ts index 08cbf391d5..f53a97bd0a 100644 --- a/packages/types/src/vscode-extension-host.ts +++ b/packages/types/src/vscode-extension-host.ts @@ -262,11 +262,13 @@ export interface ExtensionMessage { /** * CheckpointRollbackResult * - * Outcome of a change-card rollback request (B3b), posted back to the webview - * that sent `checkpointRollbackFile` / `checkpointRollbackStep`. `cardTs` - * echoes the change-card message timestamp so the requesting card can - * correlate the result: per-file results carry `filePath`, per-step results - * carry the per-file outcomes in `files`. + * Outcome of a change-card restore request (B3b), posted back to the webview + * that sent `checkpointRollbackFile` / `checkpointRollbackStep` / + * `checkpointRestoreLatestFile`. `cardTs` echoes the change-card message + * timestamp so the requesting card can correlate the result: per-file results + * carry `filePath`, per-step results carry the per-file outcomes in + * `files`, and `kind` tells a per-file result which control it belongs to + * (absent = rollback, so results posted before `kind` existed still route). */ export interface CheckpointRollbackResult { /** The `ts` of the change_card message the result belongs to. */ @@ -277,6 +279,10 @@ export interface CheckpointRollbackResult { error?: string /** Per-step scope: the per-file outcomes. */ files?: { filePath: string; success: boolean; error?: string }[] + /** Per-file scope: which control the result belongs to. Absent = rollback. */ + kind?: "rollback" | "restore-latest" + /** Per-file scope: true when a restore-latest found no recorded write and left the file as-is. */ + noOp?: boolean } export interface OpenAiCodexRateLimitsMessage { @@ -571,6 +577,7 @@ export interface WebviewMessage { | "checkpointRestore" | "checkpointRollbackFile" | "checkpointRollbackStep" + | "checkpointRestoreLatestFile" | "completionCheckpointDiff" | "completionCheckpointRestore" | "deleteMcpServer" @@ -839,6 +846,20 @@ export const checkpointRollbackStepPayloadSchema = z.object({ export type CheckpointRollbackStepPayload = z.infer +/** + * Payload of the `checkpointRestoreLatestFile` webview message (B3b): restore + * one change-card file to the latest recorded version of that file (the + * content of its most recent write checkpoint — the forward direction to a + * rollback). + */ +export const checkpointRestoreLatestFilePayloadSchema = z.object({ + /** The `ts` of the change_card message the request comes from (echoed on the result). */ + cardTs: z.number(), + filePath: z.string(), +}) + +export type CheckpointRestoreLatestFilePayload = z.infer + export interface IndexingStatusPayload { state: "Standby" | "Indexing" | "Indexed" | "Error" | "Stopping" message: string @@ -854,6 +875,7 @@ export type WebViewMessagePayload = | CheckpointRestorePayload | CheckpointRollbackFilePayload | CheckpointRollbackStepPayload + | CheckpointRestoreLatestFilePayload | IndexingStatusPayload | IndexClearedPayload | UpdateTodoListPayload diff --git a/src/core/webview/__tests__/webviewMessageHandler.rollback.spec.ts b/src/core/webview/__tests__/webviewMessageHandler.rollback.spec.ts index 74a341c26d..f75e3c4822 100644 --- a/src/core/webview/__tests__/webviewMessageHandler.rollback.spec.ts +++ b/src/core/webview/__tests__/webviewMessageHandler.rollback.spec.ts @@ -4,7 +4,7 @@ import { describe, expect, it, vi, beforeEach } from "vitest" import type { ExtensionMessage, WebviewMessage } from "@roo-code/types" import { webviewMessageHandler } from "../webviewMessageHandler" -import { rollbackFile, rollbackStep } from "../../checkpoints/rollback" +import { restoreLatestFile, rollbackFile, rollbackStep } from "../../checkpoints/rollback" import type { Task } from "../../task/Task" import type { ClineProvider } from "../ClineProvider" @@ -14,6 +14,7 @@ import type { ClineProvider } from "../ClineProvider" vi.mock("../../checkpoints/rollback", () => ({ rollbackFile: vi.fn(), rollbackStep: vi.fn(), + restoreLatestFile: vi.fn(), })) vi.mock("vscode", () => ({ @@ -197,10 +198,104 @@ describe("webviewMessageHandler - change card rollback", () => { type: "checkpointRollbackStep", payload: { cardTs: 1000, filePaths: [] }, }) + await webviewMessageHandler(provider, { + type: "checkpointRestoreLatestFile", + payload: { cardTs: 1000 } as unknown as WebviewMessage["payload"], + }) expect(rollbackFile).not.toHaveBeenCalled() expect(rollbackStep).not.toHaveBeenCalled() + expect(restoreLatestFile).not.toHaveBeenCalled() expect(postMessageToWebview).not.toHaveBeenCalled() }) }) + + describe("checkpointRestoreLatestFile", () => { + it("restores the file to its latest recorded version and posts the success outcome", async () => { + vi.mocked(restoreLatestFile).mockResolvedValueOnce({ filePath: "src/a.ts", success: true }) + + await webviewMessageHandler(provider, { + type: "checkpointRestoreLatestFile", + payload: { cardTs: 1000, filePath: "src/a.ts" }, + }) + + expect(restoreLatestFile).toHaveBeenCalledWith(mockTask, "src/a.ts") + expect(postMessageToWebview).toHaveBeenCalledWith({ + type: "checkpointRollbackResult", + checkpointRollbackResult: { + cardTs: 1000, + kind: "restore-latest", + filePath: "src/a.ts", + success: true, + }, + }) + }) + + it("flags a no-op restore-latest so the card can report it", async () => { + vi.mocked(restoreLatestFile).mockResolvedValueOnce({ filePath: "src/a.ts", success: true, noOp: true }) + + await webviewMessageHandler(provider, { + type: "checkpointRestoreLatestFile", + payload: { cardTs: 1000, filePath: "src/a.ts" }, + }) + + expect(postMessageToWebview).toHaveBeenCalledWith({ + type: "checkpointRollbackResult", + checkpointRollbackResult: { + cardTs: 1000, + kind: "restore-latest", + filePath: "src/a.ts", + success: true, + noOp: true, + }, + }) + }) + + it("posts the error outcome when the restore fails", async () => { + vi.mocked(restoreLatestFile).mockResolvedValueOnce({ + filePath: "src/a.ts", + success: false, + error: "Checkpoints are not enabled for this task", + }) + + await webviewMessageHandler(provider, { + type: "checkpointRestoreLatestFile", + payload: { cardTs: 1000, filePath: "src/a.ts" }, + }) + + expect(postMessageToWebview).toHaveBeenCalledWith({ + type: "checkpointRollbackResult", + checkpointRollbackResult: { + cardTs: 1000, + kind: "restore-latest", + filePath: "src/a.ts", + success: false, + error: "Checkpoints are not enabled for this task", + }, + }) + }) + + it("posts a correlated failure result when there is no current task", async () => { + const emptyProvider = makeProvider(undefined) + + await webviewMessageHandler(emptyProvider, { + type: "checkpointRestoreLatestFile", + payload: { cardTs: 1000, filePath: "src/a.ts" }, + }) + + expect(restoreLatestFile).not.toHaveBeenCalled() + // The requesting card must clear its pending state, so the handler + // posts a correlated failure instead of nothing. + expect(postMessageToWebview).toHaveBeenCalledWith({ + type: "checkpointRollbackResult", + checkpointRollbackResult: { + cardTs: 1000, + kind: "restore-latest", + filePath: "src/a.ts", + success: false, + error: "No active task to restore from.", + }, + }) + }) + }) }) diff --git a/src/core/webview/webviewMessageHandler.ts b/src/core/webview/webviewMessageHandler.ts index 876fec270e..5ae28bc44a 100644 --- a/src/core/webview/webviewMessageHandler.ts +++ b/src/core/webview/webviewMessageHandler.ts @@ -21,6 +21,7 @@ import { ExperimentId, checkpointRollbackFilePayloadSchema, checkpointRollbackStepPayloadSchema, + checkpointRestoreLatestFilePayloadSchema, checkoutDiffPayloadSchema, checkoutRestorePayloadSchema, getCompletionCheckpoint, @@ -1679,6 +1680,49 @@ export const webviewMessageHandler = async ( break } + case "checkpointRestoreLatestFile": { + // B3b: restore one change-card file to the latest recorded version of + // that file (the forward direction to a rollback) and report the + // outcome back to the requesting card (correlated by cardTs, kind + // "restore-latest"). + const result = checkpointRestoreLatestFilePayloadSchema.safeParse(message.payload) + + if (result.success) { + const task = provider.getCurrentTask() + + if (task) { + // Lazy import (see the checkpointRollbackFile case above). + const { restoreLatestFile } = await import("../checkpoints/rollback") + const outcome = await restoreLatestFile(task, result.data.filePath) + await provider.postMessageToWebview({ + type: "checkpointRollbackResult", + checkpointRollbackResult: { + cardTs: result.data.cardTs, + kind: "restore-latest", + filePath: outcome.filePath, + success: outcome.success, + ...(outcome.noOp ? { noOp: true } : {}), + ...(outcome.error ? { error: outcome.error } : {}), + }, + }) + } else { + // No active task: post the correlated failure so the requesting + // card can clear its pending state. + await provider.postMessageToWebview({ + type: "checkpointRollbackResult", + checkpointRollbackResult: { + cardTs: result.data.cardTs, + kind: "restore-latest", + filePath: result.data.filePath, + success: false, + error: "No active task to restore from.", + }, + }) + } + } + + break + } case "completionCheckpointDiff": { const currentCline = provider.getCurrentTask() const checkpoint = currentCline ? resolveCompletionCheckpoint(currentCline) : undefined