diff --git a/packages/types/src/vscode-extension-host.ts b/packages/types/src/vscode-extension-host.ts index b392bb77fd..08cbf391d5 100644 --- a/packages/types/src/vscode-extension-host.ts +++ b/packages/types/src/vscode-extension-host.ts @@ -108,6 +108,7 @@ export interface ExtensionMessage { | "fileContent" | "rooHistoryImportProgress" | "themeFixtureProbeRequest" + | "checkpointRollbackResult" text?: string /** For fileContent: { path, content, error? } */ fileContent?: { path: string; content: string | null; error?: string } @@ -254,6 +255,28 @@ export interface ExtensionMessage { copyProgressItemName?: string // folderSelected path?: string + /** For checkpointRollbackResult: outcome of a change-card rollback request (B3b). */ + checkpointRollbackResult?: CheckpointRollbackResult +} + +/** + * 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`. + */ +export interface CheckpointRollbackResult { + /** The `ts` of the change_card message the result belongs to. */ + cardTs: number + /** Per-file scope: the file that was restored. */ + filePath?: string + success: boolean + error?: string + /** Per-step scope: the per-file outcomes. */ + files?: { filePath: string; success: boolean; error?: string }[] } export interface OpenAiCodexRateLimitsMessage { @@ -546,6 +569,8 @@ export interface WebviewMessage { | "openCustomModesSettings" | "checkpointDiff" | "checkpointRestore" + | "checkpointRollbackFile" + | "checkpointRollbackStep" | "completionCheckpointDiff" | "completionCheckpointRestore" | "deleteMcpServer" @@ -788,6 +813,32 @@ export const checkoutRestorePayloadSchema = z.object({ export type CheckpointRestorePayload = z.infer +/** + * Payload of the `checkpointRollbackFile` webview message (B3b): restore one + * change-card file to the checkpoint commit the card was keyed by. + */ +export const checkpointRollbackFilePayloadSchema = z.object({ + /** The `ts` of the change_card message the request comes from (echoed on the result). */ + cardTs: z.number(), + checkpointId: z.string(), + filePath: z.string(), +}) + +export type CheckpointRollbackFilePayload = z.infer + +/** + * Payload of the `checkpointRollbackStep` webview message (B3b): restore + * every file of a change-card step to the step's checkpoint. + */ +export const checkpointRollbackStepPayloadSchema = z.object({ + cardTs: z.number(), + /** The step's checkpoint commit (the card's first checkpointId); optional. */ + checkpointId: z.string().optional(), + filePaths: z.array(z.string()).min(1), +}) + +export type CheckpointRollbackStepPayload = z.infer + export interface IndexingStatusPayload { state: "Standby" | "Indexing" | "Indexed" | "Error" | "Stopping" message: string @@ -801,6 +852,8 @@ export interface IndexClearedPayload { export type WebViewMessagePayload = | CheckpointDiffPayload | CheckpointRestorePayload + | CheckpointRollbackFilePayload + | CheckpointRollbackStepPayload | IndexingStatusPayload | IndexClearedPayload | UpdateTodoListPayload diff --git a/src/core/webview/__tests__/webviewMessageHandler.rollback.spec.ts b/src/core/webview/__tests__/webviewMessageHandler.rollback.spec.ts new file mode 100644 index 0000000000..74a341c26d --- /dev/null +++ b/src/core/webview/__tests__/webviewMessageHandler.rollback.spec.ts @@ -0,0 +1,206 @@ +// npx vitest run src/core/webview/__tests__/webviewMessageHandler.rollback.spec.ts +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 type { Task } from "../../task/Task" +import type { ClineProvider } from "../ClineProvider" + +// The rollback cases only call these two provider methods, so the provider +// double below is cast once at this boundary; the spy is shared so results +// can be asserted after the handler runs. +vi.mock("../../checkpoints/rollback", () => ({ + rollbackFile: vi.fn(), + rollbackStep: vi.fn(), +})) + +vi.mock("vscode", () => ({ + window: { + showErrorMessage: vi.fn(), + }, + workspace: { + workspaceFolders: undefined, + }, +})) + +// Structural mock: the handler only needs the task identity for these cases. +const mockTask = {} as Task +const postMessageToWebview = vi.fn(async (_message: ExtensionMessage) => undefined) + +function makeProvider(task: Task | undefined): ClineProvider { + const provider = { + getCurrentTask: () => task, + postMessageToWebview, + } + // Cast at the spec boundary: the rollback cases only read getCurrentTask() + // and observe postMessageToWebview calls on the structural double. + return provider as unknown as ClineProvider +} + +const provider = makeProvider(mockTask) + +describe("webviewMessageHandler - change card rollback", () => { + beforeEach(() => { + vi.clearAllMocks() + }) + + describe("checkpointRollbackFile", () => { + it("restores the file and posts the success outcome back to the webview", async () => { + vi.mocked(rollbackFile).mockResolvedValueOnce({ filePath: "src/a.ts", success: true }) + + await webviewMessageHandler(provider, { + type: "checkpointRollbackFile", + payload: { cardTs: 1000, checkpointId: "abc123", filePath: "src/a.ts" }, + }) + + expect(rollbackFile).toHaveBeenCalledWith(mockTask, "abc123", "src/a.ts") + expect(postMessageToWebview).toHaveBeenCalledWith({ + type: "checkpointRollbackResult", + checkpointRollbackResult: { cardTs: 1000, filePath: "src/a.ts", success: true }, + }) + }) + + it("posts the error outcome when the restore fails", async () => { + vi.mocked(rollbackFile).mockResolvedValueOnce({ + filePath: "src/a.ts", + success: false, + error: "checkpoint not found", + }) + + await webviewMessageHandler(provider, { + type: "checkpointRollbackFile", + payload: { cardTs: 1000, checkpointId: "abc123", filePath: "src/a.ts" }, + }) + + expect(postMessageToWebview).toHaveBeenCalledWith({ + type: "checkpointRollbackResult", + checkpointRollbackResult: { + cardTs: 1000, + filePath: "src/a.ts", + success: false, + error: "checkpoint not found", + }, + }) + }) + + it("posts a correlated failure result when there is no current task", async () => { + const emptyProvider = makeProvider(undefined) + + await webviewMessageHandler(emptyProvider, { + type: "checkpointRollbackFile", + payload: { cardTs: 1000, checkpointId: "abc123", filePath: "src/a.ts" }, + }) + + expect(rollbackFile).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, + filePath: "src/a.ts", + success: false, + error: "No active task to roll back from.", + }, + }) + }) + }) + + describe("checkpointRollbackStep", () => { + it("restores every step file and posts the aggregated outcome", async () => { + vi.mocked(rollbackStep).mockResolvedValueOnce({ + checkpointId: "abc123", + files: [ + { filePath: "src/a.ts", success: true }, + { filePath: "src/b.ts", success: true }, + ], + }) + + await webviewMessageHandler(provider, { + type: "checkpointRollbackStep", + payload: { cardTs: 1000, checkpointId: "abc123", filePaths: ["src/a.ts", "src/b.ts"] }, + }) + + expect(rollbackStep).toHaveBeenCalledWith(mockTask, ["src/a.ts", "src/b.ts"], "abc123") + expect(postMessageToWebview).toHaveBeenCalledWith({ + type: "checkpointRollbackResult", + checkpointRollbackResult: { + cardTs: 1000, + success: true, + files: [ + { filePath: "src/a.ts", success: true }, + { filePath: "src/b.ts", success: true }, + ], + }, + }) + }) + + it("reports success false with the first failing file's error when a step file fails", async () => { + vi.mocked(rollbackStep).mockResolvedValueOnce({ + checkpointId: "abc123", + files: [ + { filePath: "src/a.ts", success: true }, + { filePath: "src/b.ts", success: false, error: "boom" }, + ], + }) + + await webviewMessageHandler(provider, { + type: "checkpointRollbackStep", + payload: { cardTs: 1000, filePaths: ["src/a.ts", "src/b.ts"] }, + }) + + // Without an explicit step checkpoint id the journal lookup is used. + expect(rollbackStep).toHaveBeenCalledWith(mockTask, ["src/a.ts", "src/b.ts"], undefined) + expect(postMessageToWebview).toHaveBeenCalledWith({ + type: "checkpointRollbackResult", + checkpointRollbackResult: { + cardTs: 1000, + success: false, + error: "boom", + files: [ + { filePath: "src/a.ts", success: true }, + { filePath: "src/b.ts", success: false, error: "boom" }, + ], + }, + }) + }) + + it("posts a correlated failure result when there is no current task", async () => { + const emptyProvider = makeProvider(undefined) + + await webviewMessageHandler(emptyProvider, { + type: "checkpointRollbackStep", + payload: { cardTs: 1000, filePaths: ["src/a.ts"] }, + }) + + expect(rollbackStep).not.toHaveBeenCalled() + expect(postMessageToWebview).toHaveBeenCalledWith({ + type: "checkpointRollbackResult", + checkpointRollbackResult: { + cardTs: 1000, + success: false, + error: "No active task to roll back from.", + }, + }) + }) + + it("ignores payloads that do not match the schema", async () => { + await webviewMessageHandler(provider, { + // Malformed on purpose (only the webview produces this message): the cast + // lets the spec reach the handler's safeParse rejection without `any`. + type: "checkpointRollbackFile", + payload: { cardTs: 1000 } as unknown as WebviewMessage["payload"], + }) + await webviewMessageHandler(provider, { + type: "checkpointRollbackStep", + payload: { cardTs: 1000, filePaths: [] }, + }) + + expect(rollbackFile).not.toHaveBeenCalled() + expect(rollbackStep).not.toHaveBeenCalled() + expect(postMessageToWebview).not.toHaveBeenCalled() + }) + }) +}) diff --git a/src/core/webview/webviewMessageHandler.ts b/src/core/webview/webviewMessageHandler.ts index 0dad65a480..876fec270e 100644 --- a/src/core/webview/webviewMessageHandler.ts +++ b/src/core/webview/webviewMessageHandler.ts @@ -19,6 +19,8 @@ import { TelemetryEventName, RooCodeSettings, ExperimentId, + checkpointRollbackFilePayloadSchema, + checkpointRollbackStepPayloadSchema, checkoutDiffPayloadSchema, checkoutRestorePayloadSchema, getCompletionCheckpoint, @@ -1598,6 +1600,85 @@ export const webviewMessageHandler = async ( break } + case "checkpointRollbackFile": { + // B3b: restore one change-card file to its step checkpoint and report + // the outcome back to the requesting card (correlated by cardTs). + const result = checkpointRollbackFilePayloadSchema.safeParse(message.payload) + + if (result.success) { + const task = provider.getCurrentTask() + + if (task) { + // Lazy import: the rollback module pulls the checkpoint service and the + // editor integrations (DiffViewProvider) into the import graph. Loading it + // only when a rollback is requested keeps specs that mock `vscode` minimally + // from executing editor module-scope code at import time. + const { rollbackFile } = await import("../checkpoints/rollback") + const outcome = await rollbackFile(task, result.data.checkpointId, result.data.filePath) + await provider.postMessageToWebview({ + type: "checkpointRollbackResult", + checkpointRollbackResult: { + cardTs: result.data.cardTs, + filePath: outcome.filePath, + success: outcome.success, + ...(outcome.error ? { error: outcome.error } : {}), + }, + }) + } else { + // No active task: the rollback cannot run. Post the correlated + // failure so the requesting card can clear its pending state + // instead of waiting on a result that will never arrive. + await provider.postMessageToWebview({ + type: "checkpointRollbackResult", + checkpointRollbackResult: { + cardTs: result.data.cardTs, + filePath: result.data.filePath, + success: false, + error: "No active task to roll back from.", + }, + }) + } + } + + break + } + case "checkpointRollbackStep": { + // B3b: restore every file of a change-card step to the step checkpoint. + const result = checkpointRollbackStepPayloadSchema.safeParse(message.payload) + + if (result.success) { + const task = provider.getCurrentTask() + + if (task) { + // Lazy import (see the checkpointRollbackFile case above). + const { rollbackStep } = await import("../checkpoints/rollback") + const outcome = await rollbackStep(task, result.data.filePaths, result.data.checkpointId) + const firstFailure = outcome.files.find((file) => !file.success) + await provider.postMessageToWebview({ + type: "checkpointRollbackResult", + checkpointRollbackResult: { + cardTs: result.data.cardTs, + success: outcome.files.every((file) => file.success), + ...(firstFailure ? { error: firstFailure.error } : {}), + files: outcome.files, + }, + }) + } 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, + success: false, + error: "No active task to roll back from.", + }, + }) + } + } + + break + } case "completionCheckpointDiff": { const currentCline = provider.getCurrentTask() const checkpoint = currentCline ? resolveCompletionCheckpoint(currentCline) : undefined