Skip to content

Commit 930e70e

Browse files
liwilliam2021liwilliam2021
committed
refactor
1 parent 6044c1a commit 930e70e

File tree

3 files changed

+297
-56
lines changed

3 files changed

+297
-56
lines changed

‎src/core/webview/__tests__/checkpointRestoreHandler.test.ts‎

Lines changed: 164 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,164 @@
1+
import { describe, it, expect, vi, beforeEach } from "vitest"
2+
import { handleCheckpointRestoreOperation, hasValidCheckpoint } from "../checkpointRestoreHandler"
3+
import { saveTaskMessages } from "../../task-persistence"
4+
import * as vscode from "vscode"
5+
6+
vi.mock("../../task-persistence", () => ({
7+
saveTaskMessages: vi.fn(),
8+
}))
9+
10+
vi.mock("vscode", () => ({
11+
window: {
12+
showErrorMessage: vi.fn(),
13+
},
14+
}))
15+
16+
describe("checkpointRestoreHandler", () => {
17+
let mockProvider: any
18+
let mockCline: any
19+
20+
beforeEach(() => {
21+
vi.clearAllMocks()
22+
23+
mockProvider = {
24+
contextProxy: {
25+
globalStorageUri: {
26+
fsPath: "/test/global/storage",
27+
},
28+
},
29+
getTaskWithId: vi.fn().mockResolvedValue({
30+
historyItem: { id: "test-task", messages: [] },
31+
}),
32+
initClineWithHistoryItem: vi.fn(),
33+
}
34+
35+
mockCline = {
36+
taskId: "test-task",
37+
clineMessages: [
38+
{ ts: 1, text: "Message 1" },
39+
{ ts: 2, text: "Message 2", checkpoint: { hash: "abc123" } },
40+
{ ts: 3, text: "Message 3" },
41+
],
42+
checkpointRestore: vi.fn(),
43+
}
44+
})
45+
46+
describe("hasValidCheckpoint", () => {
47+
it("should return true for valid checkpoint", () => {
48+
const message = { checkpoint: { hash: "abc123" } }
49+
expect(hasValidCheckpoint(message)).toBe(true)
50+
})
51+
52+
it("should return false for missing checkpoint", () => {
53+
const message = { text: "No checkpoint" }
54+
expect(hasValidCheckpoint(message)).toBe(false)
55+
})
56+
57+
it("should return false for invalid checkpoint structure", () => {
58+
expect(hasValidCheckpoint({ checkpoint: "invalid" })).toBe(false)
59+
expect(hasValidCheckpoint({ checkpoint: {} })).toBe(false)
60+
expect(hasValidCheckpoint({ checkpoint: { hash: 123 } })).toBe(false)
61+
})
62+
})
63+
64+
describe("handleCheckpointRestoreOperation", () => {
65+
describe("delete operation", () => {
66+
it("should handle delete operation correctly", async () => {
67+
await handleCheckpointRestoreOperation({
68+
provider: mockProvider,
69+
currentCline: mockCline,
70+
messageTs: 2,
71+
messageIndex: 1,
72+
checkpoint: { hash: "abc123" },
73+
operation: "delete",
74+
})
75+
76+
// Should call checkpointRestore with correct params
77+
expect(mockCline.checkpointRestore).toHaveBeenCalledWith({
78+
ts: 2,
79+
commitHash: "abc123",
80+
mode: "restore",
81+
operation: "delete",
82+
})
83+
84+
// Should save messages after restoration
85+
expect(saveTaskMessages).toHaveBeenCalledWith({
86+
messages: mockCline.clineMessages,
87+
taskId: "test-task",
88+
globalStoragePath: "/test/global/storage",
89+
})
90+
91+
// Should reinitialize the task
92+
expect(mockProvider.getTaskWithId).toHaveBeenCalledWith("test-task")
93+
expect(mockProvider.initClineWithHistoryItem).toHaveBeenCalledWith({
94+
id: "test-task",
95+
messages: [],
96+
})
97+
})
98+
})
99+
100+
describe("edit operation", () => {
101+
it("should handle edit operation correctly", async () => {
102+
const editData = {
103+
editedContent: "Edited content",
104+
images: ["image1.png"],
105+
apiConversationHistoryIndex: 1,
106+
}
107+
108+
await handleCheckpointRestoreOperation({
109+
provider: mockProvider,
110+
currentCline: mockCline,
111+
messageTs: 2,
112+
messageIndex: 1,
113+
checkpoint: { hash: "abc123" },
114+
operation: "edit",
115+
editData,
116+
})
117+
118+
// Should set pendingEditAfterRestore on provider
119+
expect(mockProvider.pendingEditAfterRestore).toEqual({
120+
messageTs: 2,
121+
editedContent: "Edited content",
122+
images: ["image1.png"],
123+
messageIndex: 1,
124+
apiConversationHistoryIndex: 1,
125+
originalCheckpoint: { hash: "abc123" },
126+
})
127+
128+
// Should call checkpointRestore with correct params
129+
expect(mockCline.checkpointRestore).toHaveBeenCalledWith({
130+
ts: 2,
131+
commitHash: "abc123",
132+
mode: "restore",
133+
operation: "edit",
134+
})
135+
136+
// Should NOT save messages or reinitialize for edit
137+
expect(saveTaskMessages).not.toHaveBeenCalled()
138+
expect(mockProvider.initClineWithHistoryItem).not.toHaveBeenCalled()
139+
})
140+
})
141+
142+
describe("error handling", () => {
143+
it("should handle errors and show error message", async () => {
144+
const error = new Error("Checkpoint restore failed")
145+
mockCline.checkpointRestore.mockRejectedValue(error)
146+
147+
await expect(
148+
handleCheckpointRestoreOperation({
149+
provider: mockProvider,
150+
currentCline: mockCline,
151+
messageTs: 2,
152+
messageIndex: 1,
153+
checkpoint: { hash: "abc123" },
154+
operation: "delete",
155+
}),
156+
).rejects.toThrow("Checkpoint restore failed")
157+
158+
expect(vscode.window.showErrorMessage).toHaveBeenCalledWith(
159+
"Error during checkpoint restore: Checkpoint restore failed",
160+
)
161+
})
162+
})
163+
})
164+
})

‎src/core/webview/checkpointRestoreHandler.ts‎

Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,103 @@
1+
import { Task } from "../task/Task"
2+
import { ClineProvider } from "./ClineProvider"
3+
import { saveTaskMessages } from "../task-persistence"
4+
import * as vscode from "vscode"
5+
import pWaitFor from "p-wait-for"
6+
import { t } from "../../i18n"
7+
8+
export interface CheckpointRestoreConfig {
9+
provider: ClineProvider
10+
currentCline: Task
11+
messageTs: number
12+
messageIndex: number
13+
checkpoint: { hash: string }
14+
operation: "delete" | "edit"
15+
editData?: {
16+
editedContent: string
17+
images?: string[]
18+
apiConversationHistoryIndex: number
19+
}
20+
}
21+
22+
/**
23+
* Handles checkpoint restoration for both delete and edit operations.
24+
* This consolidates the common logic while handling operation-specific behavior.
25+
*/
26+
export async function handleCheckpointRestoreOperation(config: CheckpointRestoreConfig): Promise<void> {
27+
const { provider, currentCline, messageTs, checkpoint, operation, editData } = config
28+
29+
try {
30+
// For edit operations, set up pending edit data before restoration
31+
if (operation === "edit" && editData) {
32+
;(provider as any).pendingEditAfterRestore = {
33+
messageTs,
34+
editedContent: editData.editedContent,
35+
images: editData.images,
36+
messageIndex: config.messageIndex,
37+
apiConversationHistoryIndex: editData.apiConversationHistoryIndex,
38+
originalCheckpoint: checkpoint,
39+
}
40+
}
41+
42+
// Perform the checkpoint restoration
43+
await currentCline.checkpointRestore({
44+
ts: messageTs,
45+
commitHash: checkpoint.hash,
46+
mode: "restore",
47+
operation,
48+
})
49+
50+
// For delete operations, we need to save messages and reinitialize
51+
// For edit operations, the reinitialization happens automatically
52+
// and processes the pending edit
53+
if (operation === "delete") {
54+
// Save the updated messages to disk after checkpoint restoration
55+
await saveTaskMessages({
56+
messages: currentCline.clineMessages,
57+
taskId: currentCline.taskId,
58+
globalStoragePath: provider.contextProxy.globalStorageUri.fsPath,
59+
})
60+
61+
// Get the updated history item and reinitialize
62+
const { historyItem } = await provider.getTaskWithId(currentCline.taskId)
63+
await provider.initClineWithHistoryItem(historyItem)
64+
}
65+
// For edit operations, the task cancellation in checkpointRestore
66+
// will trigger reinitialization, which will process pendingEditAfterRestore
67+
} catch (error) {
68+
console.error(`Error in checkpoint restore (${operation}):`, error)
69+
vscode.window.showErrorMessage(
70+
`Error during checkpoint restore: ${error instanceof Error ? error.message : String(error)}`,
71+
)
72+
throw error
73+
}
74+
}
75+
76+
/**
77+
* Validates if a message has a valid checkpoint for restoration
78+
*/
79+
export function hasValidCheckpoint(message: any): boolean {
80+
return (
81+
(message?.checkpoint &&
82+
typeof message.checkpoint === "object" &&
83+
"hash" in message.checkpoint &&
84+
typeof message.checkpoint.hash === "string") ||
85+
false
86+
)
87+
}
88+
89+
/**
90+
* Common checkpoint restore validation and initialization utility.
91+
* This can be used by any checkpoint restore flow that needs to wait for initialization.
92+
*/
93+
export async function waitForClineInitialization(provider: ClineProvider, timeoutMs: number = 3000): Promise<boolean> {
94+
try {
95+
await pWaitFor(() => provider.getCurrentCline()?.isInitialized === true, {
96+
timeout: timeoutMs,
97+
})
98+
return true
99+
} catch (error) {
100+
vscode.window.showErrorMessage(t("common:errors.checkpoint_timeout"))
101+
return false
102+
}
103+
}

‎src/core/webview/webviewMessageHandler.ts‎

Lines changed: 30 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,11 @@ import { type ApiMessage } from "../task-persistence/apiMessages"
1919
import { saveTaskMessages } from "../task-persistence"
2020

2121
import { ClineProvider } from "./ClineProvider"
22+
import {
23+
handleCheckpointRestoreOperation,
24+
hasValidCheckpoint,
25+
waitForClineInitialization,
26+
} from "./checkpointRestoreHandler"
2227
import { changeLanguage, t } from "../../i18n"
2328
import { Package } from "../../shared/package"
2429
import { RouterName, toRouterName, ModelRecord } from "../../shared/api"
@@ -154,29 +159,18 @@ export const webviewMessageHandler = async (
154159

155160
if (messageIndex !== -1) {
156161
try {
157-
// If checkpoint restoration is requested, restore to the checkpoint first
158-
if (restoreCheckpoint) {
159-
const targetMessage = currentCline.clineMessages[messageIndex]
160-
if (
161-
targetMessage?.checkpoint &&
162-
typeof targetMessage.checkpoint === "object" &&
163-
"hash" in targetMessage.checkpoint
164-
) {
165-
await currentCline.checkpointRestore({
166-
ts: targetMessage.ts!,
167-
commitHash: targetMessage.checkpoint.hash as string,
168-
mode: "restore",
169-
operation: "delete",
170-
})
162+
const targetMessage = currentCline.clineMessages[messageIndex]
171163

172-
// Save the updated messages to disk after checkpoint restoration
173-
// This ensures the deleted messages are persisted before reinitialization
174-
await saveTaskMessages({
175-
messages: currentCline.clineMessages,
176-
taskId: currentCline.taskId,
177-
globalStoragePath: provider.contextProxy.globalStorageUri.fsPath,
178-
})
179-
}
164+
// If checkpoint restoration is requested, restore to the checkpoint first
165+
if (restoreCheckpoint && hasValidCheckpoint(targetMessage)) {
166+
await handleCheckpointRestoreOperation({
167+
provider,
168+
currentCline,
169+
messageTs: targetMessage.ts!,
170+
messageIndex,
171+
checkpoint: targetMessage.checkpoint as { hash: string },
172+
operation: "delete",
173+
})
180174
} else {
181175
// For non-checkpoint deletes, preserve checkpoint associations for remaining messages
182176
// Store checkpoints from messages that will be preserved
@@ -206,13 +200,6 @@ export const webviewMessageHandler = async (
206200
globalStoragePath: provider.contextProxy.globalStorageUri.fsPath,
207201
})
208202
}
209-
210-
const { historyItem } = await provider.getTaskWithId(currentCline.taskId)
211-
212-
// Initialize with history item after deletion (only for checkpoint restores)
213-
if (restoreCheckpoint) {
214-
await provider.initClineWithHistoryItem(historyItem)
215-
}
216203
} catch (error) {
217204
console.error("Error in delete message:", error)
218205
vscode.window.showErrorMessage(
@@ -307,36 +294,23 @@ export const webviewMessageHandler = async (
307294
const originalCheckpoint = targetMessage?.checkpoint
308295

309296
// If checkpoint restoration is requested, restore to the checkpoint first
310-
if (restoreCheckpoint) {
311-
if (
312-
originalCheckpoint &&
313-
typeof originalCheckpoint === "object" &&
314-
"hash" in originalCheckpoint
315-
) {
316-
// Store the edited content and images for after restoration
317-
const editData = { text: editedContent, images }
318-
319-
// Set a flag on the provider to indicate we need to process an edit after restoration
320-
;(provider as any).pendingEditAfterRestore = {
321-
messageTs,
297+
if (restoreCheckpoint && hasValidCheckpoint(targetMessage)) {
298+
await handleCheckpointRestoreOperation({
299+
provider,
300+
currentCline,
301+
messageTs: targetMessage.ts!,
302+
messageIndex,
303+
checkpoint: targetMessage.checkpoint as { hash: string },
304+
operation: "edit",
305+
editData: {
322306
editedContent,
323307
images,
324-
messageIndex,
325308
apiConversationHistoryIndex,
326-
originalCheckpoint, // Preserve the checkpoint for the new message
327-
}
328-
329-
await currentCline.checkpointRestore({
330-
ts: targetMessage.ts!,
331-
commitHash: originalCheckpoint.hash as string,
332-
mode: "restore",
333-
operation: "edit",
334-
})
335-
336-
// The task will be cancelled and reinitialized by checkpointRestore
337-
// The pending edit will be processed in the reinitialized task
338-
return
339-
}
309+
},
310+
})
311+
// The task will be cancelled and reinitialized by checkpointRestore
312+
// The pending edit will be processed in the reinitialized task
313+
return
340314
}
341315

342316
// For non-checkpoint edits, preserve checkpoint associations for remaining messages

0 commit comments

Comments
 (0)