Skip to content

Commit a0bb49b

Browse files
easonliang28easonliang28
committed
feat(webview): change cards UI and rollback buttons (B3b, #1375)
1 parent 2b4a8ce commit a0bb49b

28 files changed

Lines changed: 1300 additions & 1 deletion

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

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -108,6 +108,7 @@ export interface ExtensionMessage {
108108
| "fileContent"
109109
| "rooHistoryImportProgress"
110110
| "themeFixtureProbeRequest"
111+
| "checkpointRollbackResult"
111112
text?: string
112113
/** For fileContent: { path, content, error? } */
113114
fileContent?: { path: string; content: string | null; error?: string }
@@ -254,6 +255,28 @@ export interface ExtensionMessage {
254255
copyProgressItemName?: string
255256
// folderSelected
256257
path?: string
258+
/** For checkpointRollbackResult: outcome of a change-card rollback request (B3b). */
259+
checkpointRollbackResult?: CheckpointRollbackResult
260+
}
261+
262+
/**
263+
* CheckpointRollbackResult
264+
*
265+
* Outcome of a change-card rollback request (B3b), posted back to the webview
266+
* that sent `checkpointRollbackFile` / `checkpointRollbackStep`. `cardTs`
267+
* echoes the change-card message timestamp so the requesting card can
268+
* correlate the result: per-file results carry `filePath`, per-step results
269+
* carry the per-file outcomes in `files`.
270+
*/
271+
export interface CheckpointRollbackResult {
272+
/** The `ts` of the change_card message the result belongs to. */
273+
cardTs: number
274+
/** Per-file scope: the file that was restored. */
275+
filePath?: string
276+
success: boolean
277+
error?: string
278+
/** Per-step scope: the per-file outcomes. */
279+
files?: { filePath: string; success: boolean; error?: string }[]
257280
}
258281

259282
export interface OpenAiCodexRateLimitsMessage {
@@ -546,6 +569,8 @@ export interface WebviewMessage {
546569
| "openCustomModesSettings"
547570
| "checkpointDiff"
548571
| "checkpointRestore"
572+
| "checkpointRollbackFile"
573+
| "checkpointRollbackStep"
549574
| "completionCheckpointDiff"
550575
| "completionCheckpointRestore"
551576
| "deleteMcpServer"
@@ -788,6 +813,32 @@ export const checkoutRestorePayloadSchema = z.object({
788813

789814
export type CheckpointRestorePayload = z.infer<typeof checkoutRestorePayloadSchema>
790815

816+
/**
817+
* Payload of the `checkpointRollbackFile` webview message (B3b): restore one
818+
* change-card file to the checkpoint commit the card was keyed by.
819+
*/
820+
export const checkpointRollbackFilePayloadSchema = z.object({
821+
/** The `ts` of the change_card message the request comes from (echoed on the result). */
822+
cardTs: z.number(),
823+
checkpointId: z.string(),
824+
filePath: z.string(),
825+
})
826+
827+
export type CheckpointRollbackFilePayload = z.infer<typeof checkpointRollbackFilePayloadSchema>
828+
829+
/**
830+
* Payload of the `checkpointRollbackStep` webview message (B3b): restore
831+
* every file of a change-card step to the step's checkpoint.
832+
*/
833+
export const checkpointRollbackStepPayloadSchema = z.object({
834+
cardTs: z.number(),
835+
/** The step's checkpoint commit (the card's first checkpointId); optional. */
836+
checkpointId: z.string().optional(),
837+
filePaths: z.array(z.string()).min(1),
838+
})
839+
840+
export type CheckpointRollbackStepPayload = z.infer<typeof checkpointRollbackStepPayloadSchema>
841+
791842
export interface IndexingStatusPayload {
792843
state: "Standby" | "Indexing" | "Indexed" | "Error" | "Stopping"
793844
message: string
@@ -801,6 +852,8 @@ export interface IndexClearedPayload {
801852
export type WebViewMessagePayload =
802853
| CheckpointDiffPayload
803854
| CheckpointRestorePayload
855+
| CheckpointRollbackFilePayload
856+
| CheckpointRollbackStepPayload
804857
| IndexingStatusPayload
805858
| IndexClearedPayload
806859
| UpdateTodoListPayload

‎src/core/webview/__tests__/webviewMessageHandler.rollback.spec.ts‎

Lines changed: 206 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,206 @@
1+
// npx vitest run src/core/webview/__tests__/webviewMessageHandler.rollback.spec.ts
2+
import { describe, expect, it, vi, beforeEach } from "vitest"
3+
4+
import type { ExtensionMessage, WebviewMessage } from "@roo-code/types"
5+
6+
import { webviewMessageHandler } from "../webviewMessageHandler"
7+
import { rollbackFile, rollbackStep } from "../../checkpoints/rollback"
8+
import type { Task } from "../../task/Task"
9+
import type { ClineProvider } from "../ClineProvider"
10+
11+
// The rollback cases only call these two provider methods, so the provider
12+
// double below is cast once at this boundary; the spy is shared so results
13+
// can be asserted after the handler runs.
14+
vi.mock("../../checkpoints/rollback", () => ({
15+
rollbackFile: vi.fn(),
16+
rollbackStep: vi.fn(),
17+
}))
18+
19+
vi.mock("vscode", () => ({
20+
window: {
21+
showErrorMessage: vi.fn(),
22+
},
23+
workspace: {
24+
workspaceFolders: undefined,
25+
},
26+
}))
27+
28+
// Structural mock: the handler only needs the task identity for these cases.
29+
const mockTask = {} as Task
30+
const postMessageToWebview = vi.fn(async (_message: ExtensionMessage) => undefined)
31+
32+
function makeProvider(task: Task | undefined): ClineProvider {
33+
const provider = {
34+
getCurrentTask: () => task,
35+
postMessageToWebview,
36+
}
37+
// Cast at the spec boundary: the rollback cases only read getCurrentTask()
38+
// and observe postMessageToWebview calls on the structural double.
39+
return provider as unknown as ClineProvider
40+
}
41+
42+
const provider = makeProvider(mockTask)
43+
44+
describe("webviewMessageHandler - change card rollback", () => {
45+
beforeEach(() => {
46+
vi.clearAllMocks()
47+
})
48+
49+
describe("checkpointRollbackFile", () => {
50+
it("restores the file and posts the success outcome back to the webview", async () => {
51+
vi.mocked(rollbackFile).mockResolvedValueOnce({ filePath: "src/a.ts", success: true })
52+
53+
await webviewMessageHandler(provider, {
54+
type: "checkpointRollbackFile",
55+
payload: { cardTs: 1000, checkpointId: "abc123", filePath: "src/a.ts" },
56+
})
57+
58+
expect(rollbackFile).toHaveBeenCalledWith(mockTask, "abc123", "src/a.ts")
59+
expect(postMessageToWebview).toHaveBeenCalledWith({
60+
type: "checkpointRollbackResult",
61+
checkpointRollbackResult: { cardTs: 1000, filePath: "src/a.ts", success: true },
62+
})
63+
})
64+
65+
it("posts the error outcome when the restore fails", async () => {
66+
vi.mocked(rollbackFile).mockResolvedValueOnce({
67+
filePath: "src/a.ts",
68+
success: false,
69+
error: "checkpoint not found",
70+
})
71+
72+
await webviewMessageHandler(provider, {
73+
type: "checkpointRollbackFile",
74+
payload: { cardTs: 1000, checkpointId: "abc123", filePath: "src/a.ts" },
75+
})
76+
77+
expect(postMessageToWebview).toHaveBeenCalledWith({
78+
type: "checkpointRollbackResult",
79+
checkpointRollbackResult: {
80+
cardTs: 1000,
81+
filePath: "src/a.ts",
82+
success: false,
83+
error: "checkpoint not found",
84+
},
85+
})
86+
})
87+
88+
it("posts a correlated failure result when there is no current task", async () => {
89+
const emptyProvider = makeProvider(undefined)
90+
91+
await webviewMessageHandler(emptyProvider, {
92+
type: "checkpointRollbackFile",
93+
payload: { cardTs: 1000, checkpointId: "abc123", filePath: "src/a.ts" },
94+
})
95+
96+
expect(rollbackFile).not.toHaveBeenCalled()
97+
// The requesting card must clear its pending state, so the handler
98+
// posts a correlated failure instead of nothing.
99+
expect(postMessageToWebview).toHaveBeenCalledWith({
100+
type: "checkpointRollbackResult",
101+
checkpointRollbackResult: {
102+
cardTs: 1000,
103+
filePath: "src/a.ts",
104+
success: false,
105+
error: "No active task to roll back from.",
106+
},
107+
})
108+
})
109+
})
110+
111+
describe("checkpointRollbackStep", () => {
112+
it("restores every step file and posts the aggregated outcome", async () => {
113+
vi.mocked(rollbackStep).mockResolvedValueOnce({
114+
checkpointId: "abc123",
115+
files: [
116+
{ filePath: "src/a.ts", success: true },
117+
{ filePath: "src/b.ts", success: true },
118+
],
119+
})
120+
121+
await webviewMessageHandler(provider, {
122+
type: "checkpointRollbackStep",
123+
payload: { cardTs: 1000, checkpointId: "abc123", filePaths: ["src/a.ts", "src/b.ts"] },
124+
})
125+
126+
expect(rollbackStep).toHaveBeenCalledWith(mockTask, ["src/a.ts", "src/b.ts"], "abc123")
127+
expect(postMessageToWebview).toHaveBeenCalledWith({
128+
type: "checkpointRollbackResult",
129+
checkpointRollbackResult: {
130+
cardTs: 1000,
131+
success: true,
132+
files: [
133+
{ filePath: "src/a.ts", success: true },
134+
{ filePath: "src/b.ts", success: true },
135+
],
136+
},
137+
})
138+
})
139+
140+
it("reports success false with the first failing file's error when a step file fails", async () => {
141+
vi.mocked(rollbackStep).mockResolvedValueOnce({
142+
checkpointId: "abc123",
143+
files: [
144+
{ filePath: "src/a.ts", success: true },
145+
{ filePath: "src/b.ts", success: false, error: "boom" },
146+
],
147+
})
148+
149+
await webviewMessageHandler(provider, {
150+
type: "checkpointRollbackStep",
151+
payload: { cardTs: 1000, filePaths: ["src/a.ts", "src/b.ts"] },
152+
})
153+
154+
// Without an explicit step checkpoint id the journal lookup is used.
155+
expect(rollbackStep).toHaveBeenCalledWith(mockTask, ["src/a.ts", "src/b.ts"], undefined)
156+
expect(postMessageToWebview).toHaveBeenCalledWith({
157+
type: "checkpointRollbackResult",
158+
checkpointRollbackResult: {
159+
cardTs: 1000,
160+
success: false,
161+
error: "boom",
162+
files: [
163+
{ filePath: "src/a.ts", success: true },
164+
{ filePath: "src/b.ts", success: false, error: "boom" },
165+
],
166+
},
167+
})
168+
})
169+
170+
it("posts a correlated failure result when there is no current task", async () => {
171+
const emptyProvider = makeProvider(undefined)
172+
173+
await webviewMessageHandler(emptyProvider, {
174+
type: "checkpointRollbackStep",
175+
payload: { cardTs: 1000, filePaths: ["src/a.ts"] },
176+
})
177+
178+
expect(rollbackStep).not.toHaveBeenCalled()
179+
expect(postMessageToWebview).toHaveBeenCalledWith({
180+
type: "checkpointRollbackResult",
181+
checkpointRollbackResult: {
182+
cardTs: 1000,
183+
success: false,
184+
error: "No active task to roll back from.",
185+
},
186+
})
187+
})
188+
189+
it("ignores payloads that do not match the schema", async () => {
190+
await webviewMessageHandler(provider, {
191+
// Malformed on purpose (only the webview produces this message): the cast
192+
// lets the spec reach the handler's safeParse rejection without `any`.
193+
type: "checkpointRollbackFile",
194+
payload: { cardTs: 1000 } as unknown as WebviewMessage["payload"],
195+
})
196+
await webviewMessageHandler(provider, {
197+
type: "checkpointRollbackStep",
198+
payload: { cardTs: 1000, filePaths: [] },
199+
})
200+
201+
expect(rollbackFile).not.toHaveBeenCalled()
202+
expect(rollbackStep).not.toHaveBeenCalled()
203+
expect(postMessageToWebview).not.toHaveBeenCalled()
204+
})
205+
})
206+
})

‎src/core/webview/webviewMessageHandler.ts‎

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,8 @@ import {
1919
TelemetryEventName,
2020
RooCodeSettings,
2121
ExperimentId,
22+
checkpointRollbackFilePayloadSchema,
23+
checkpointRollbackStepPayloadSchema,
2224
checkoutDiffPayloadSchema,
2325
checkoutRestorePayloadSchema,
2426
getCompletionCheckpoint,
@@ -1598,6 +1600,85 @@ export const webviewMessageHandler = async (
15981600

15991601
break
16001602
}
1603+
case "checkpointRollbackFile": {
1604+
// B3b: restore one change-card file to its step checkpoint and report
1605+
// the outcome back to the requesting card (correlated by cardTs).
1606+
const result = checkpointRollbackFilePayloadSchema.safeParse(message.payload)
1607+
1608+
if (result.success) {
1609+
const task = provider.getCurrentTask()
1610+
1611+
if (task) {
1612+
// Lazy import: the rollback module pulls the checkpoint service and the
1613+
// editor integrations (DiffViewProvider) into the import graph. Loading it
1614+
// only when a rollback is requested keeps specs that mock `vscode` minimally
1615+
// from executing editor module-scope code at import time.
1616+
const { rollbackFile } = await import("../checkpoints/rollback")
1617+
const outcome = await rollbackFile(task, result.data.checkpointId, result.data.filePath)
1618+
await provider.postMessageToWebview({
1619+
type: "checkpointRollbackResult",
1620+
checkpointRollbackResult: {
1621+
cardTs: result.data.cardTs,
1622+
filePath: outcome.filePath,
1623+
success: outcome.success,
1624+
...(outcome.error ? { error: outcome.error } : {}),
1625+
},
1626+
})
1627+
} else {
1628+
// No active task: the rollback cannot run. Post the correlated
1629+
// failure so the requesting card can clear its pending state
1630+
// instead of waiting on a result that will never arrive.
1631+
await provider.postMessageToWebview({
1632+
type: "checkpointRollbackResult",
1633+
checkpointRollbackResult: {
1634+
cardTs: result.data.cardTs,
1635+
filePath: result.data.filePath,
1636+
success: false,
1637+
error: "No active task to roll back from.",
1638+
},
1639+
})
1640+
}
1641+
}
1642+
1643+
break
1644+
}
1645+
case "checkpointRollbackStep": {
1646+
// B3b: restore every file of a change-card step to the step checkpoint.
1647+
const result = checkpointRollbackStepPayloadSchema.safeParse(message.payload)
1648+
1649+
if (result.success) {
1650+
const task = provider.getCurrentTask()
1651+
1652+
if (task) {
1653+
// Lazy import (see the checkpointRollbackFile case above).
1654+
const { rollbackStep } = await import("../checkpoints/rollback")
1655+
const outcome = await rollbackStep(task, result.data.filePaths, result.data.checkpointId)
1656+
const firstFailure = outcome.files.find((file) => !file.success)
1657+
await provider.postMessageToWebview({
1658+
type: "checkpointRollbackResult",
1659+
checkpointRollbackResult: {
1660+
cardTs: result.data.cardTs,
1661+
success: outcome.files.every((file) => file.success),
1662+
...(firstFailure ? { error: firstFailure.error } : {}),
1663+
files: outcome.files,
1664+
},
1665+
})
1666+
} else {
1667+
// No active task: post the correlated failure so the requesting
1668+
// card can clear its pending state.
1669+
await provider.postMessageToWebview({
1670+
type: "checkpointRollbackResult",
1671+
checkpointRollbackResult: {
1672+
cardTs: result.data.cardTs,
1673+
success: false,
1674+
error: "No active task to roll back from.",
1675+
},
1676+
})
1677+
}
1678+
}
1679+
1680+
break
1681+
}
16011682
case "completionCheckpointDiff": {
16021683
const currentCline = provider.getCurrentTask()
16031684
const checkpoint = currentCline ? resolveCompletionCheckpoint(currentCline) : undefined

0 commit comments

Comments
 (0)