Skip to content

Commit cfae65f

Browse files
ivanarifinivanarifin
committed
fix: address completion action review feedback
1 parent 54ee431 commit cfae65f

12 files changed

Lines changed: 170 additions & 68 deletions

‎packages/types/src/__tests__/message.test.ts‎

Lines changed: 52 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,14 @@
11
// pnpm --filter @roo-code/types test src/__tests__/message.test.ts
22

3-
import { clineAsks, isIdleAsk, isInteractiveAsk, isResumableAsk, isNonBlockingAsk } from "../message.js"
3+
import {
4+
clineAsks,
5+
getCompletionCheckpoint,
6+
isIdleAsk,
7+
isInteractiveAsk,
8+
isResumableAsk,
9+
isNonBlockingAsk,
10+
type ClineMessage,
11+
} from "../message.js"
412

513
describe("ask messages", () => {
614
test("all ask messages are classified", () => {
@@ -12,3 +20,46 @@ describe("ask messages", () => {
1220
}
1321
})
1422
})
23+
24+
describe("getCompletionCheckpoint", () => {
25+
it("returns the first checkpoint after the latest user prompt before completion", () => {
26+
const messages: ClineMessage[] = [
27+
{ type: "say", say: "text", ts: 1, text: "Initial task" },
28+
{ type: "say", say: "checkpoint_saved", ts: 2, text: "initial-checkpoint" },
29+
{ type: "say", say: "completion_result", ts: 3, text: "First completion" },
30+
{ type: "say", say: "user_feedback", ts: 4, text: "Change it" },
31+
{ type: "say", say: "checkpoint_saved", ts: 5, text: "latest-prompt-checkpoint" },
32+
{ type: "say", say: "checkpoint_saved", ts: 6, text: "later-edit-checkpoint" },
33+
{ type: "ask", ask: "completion_result", ts: 7, text: "", partial: false },
34+
]
35+
36+
expect(getCompletionCheckpoint(messages)).toEqual({
37+
ts: 5,
38+
commitHash: "latest-prompt-checkpoint",
39+
})
40+
})
41+
42+
it("returns the first checkpoint after an initial task row before completion", () => {
43+
const messages = [
44+
{ type: "say", say: "task", ts: 1, text: "Initial task" },
45+
{ type: "say", say: "checkpoint_saved", ts: 2, text: "checkpoint-after-initial-task" },
46+
{ type: "ask", ask: "completion_result", ts: 3, text: "Task complete", partial: false },
47+
] as ClineMessage[]
48+
49+
expect(getCompletionCheckpoint(messages)).toEqual({
50+
ts: 2,
51+
commitHash: "checkpoint-after-initial-task",
52+
})
53+
})
54+
55+
it("returns undefined when completion has no checkpoint after the latest user prompt", () => {
56+
const messages: ClineMessage[] = [
57+
{ type: "say", say: "text", ts: 1, text: "Initial task" },
58+
{ type: "say", say: "checkpoint_saved", ts: 2, text: "initial-checkpoint" },
59+
{ type: "say", say: "user_feedback", ts: 3, text: "Change it" },
60+
{ type: "ask", ask: "completion_result", ts: 4, text: "", partial: false },
61+
]
62+
63+
expect(getCompletionCheckpoint(messages)).toBeUndefined()
64+
})
65+
})

‎packages/types/src/message.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -281,7 +281,7 @@ export interface CompletionCheckpoint {
281281
}
282282

283283
const isInitialTaskMessage = (message: ClineMessage | undefined): boolean => {
284-
return message?.type === "say" && message.say === "text"
284+
return message?.type === "say" && (message.say === "text" || (message as { say?: string }).say === "task")
285285
}
286286

287287
const isUserFeedbackMessage = (message: ClineMessage): boolean => {

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -645,6 +645,7 @@ export interface WebviewMessage {
645645
ids?: string[]
646646
terminalOperation?: "continue" | "abort"
647647
messageTs?: number
648+
checkpointTs?: number
648649
restoreCheckpoint?: boolean
649650
historyPreviewCollapsed?: boolean
650651
filters?: { type?: string; search?: string; tags?: string[] }

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

Lines changed: 0 additions & 31 deletions
This file was deleted.

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

Lines changed: 28 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import { describe, it, expect, vi, beforeEach } from "vitest"
2+
import pWaitFor from "p-wait-for"
23
import { webviewMessageHandler } from "../webviewMessageHandler"
34
import { saveTaskMessages } from "../../task-persistence"
45
import { handleCheckpointRestoreOperation } from "../checkpointRestoreHandler"
@@ -7,6 +8,13 @@ import { MessageManager } from "../../message-manager"
78
// Mock dependencies
89
vi.mock("../../task-persistence")
910
vi.mock("../checkpointRestoreHandler")
11+
vi.mock("p-wait-for", () => ({
12+
default: vi.fn(async (condition: () => boolean) => {
13+
if (!condition()) {
14+
throw new Error("condition not met")
15+
}
16+
}),
17+
}))
1018
vi.mock("vscode", () => ({
1119
window: {
1220
showErrorMessage: vi.fn(),
@@ -151,7 +159,7 @@ describe("webviewMessageHandler - checkpoint operations", () => {
151159
})
152160

153161
it("diffs changes from the checkpoint created after the latest prompt", async () => {
154-
await webviewMessageHandler(mockProvider, { type: "completionCheckpointDiff" })
162+
await webviewMessageHandler(mockProvider, { type: "completionCheckpointDiff", checkpointTs: 4 })
155163

156164
expect(mockCline.checkpointDiff).toHaveBeenCalledWith({
157165
ts: 4,
@@ -161,7 +169,7 @@ describe("webviewMessageHandler - checkpoint operations", () => {
161169
})
162170

163171
it("restores files and task state to the checkpoint created after the latest prompt", async () => {
164-
await webviewMessageHandler(mockProvider, { type: "completionCheckpointRestore" })
172+
await webviewMessageHandler(mockProvider, { type: "completionCheckpointRestore", checkpointTs: 4 })
165173

166174
expect(mockProvider.cancelTask).toHaveBeenCalled()
167175
expect(mockCline.checkpointRestore).toHaveBeenCalledWith({
@@ -185,5 +193,23 @@ describe("webviewMessageHandler - checkpoint operations", () => {
185193
expect(mockCline.checkpointRestore).not.toHaveBeenCalled()
186194
expect(mockProvider.cancelTask).not.toHaveBeenCalled()
187195
})
196+
197+
it("does not re-derive a different checkpoint when an explicit checkpoint timestamp is missing", async () => {
198+
await webviewMessageHandler(mockProvider, { type: "completionCheckpointDiff", checkpointTs: 999 })
199+
await webviewMessageHandler(mockProvider, { type: "completionCheckpointRestore", checkpointTs: 999 })
200+
201+
expect(mockCline.checkpointDiff).not.toHaveBeenCalled()
202+
expect(mockCline.checkpointRestore).not.toHaveBeenCalled()
203+
expect(mockProvider.cancelTask).not.toHaveBeenCalled()
204+
})
205+
206+
it("does not restore when task re-initialization times out", async () => {
207+
;(pWaitFor as any).mockRejectedValueOnce(new Error("timed out"))
208+
209+
await webviewMessageHandler(mockProvider, { type: "completionCheckpointRestore", checkpointTs: 4 })
210+
211+
expect(mockProvider.cancelTask).toHaveBeenCalled()
212+
expect(mockCline.checkpointRestore).not.toHaveBeenCalled()
213+
})
188214
})
189215
})

‎src/core/webview/webviewMessageHandler.ts‎

Lines changed: 26 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,25 @@ export const webviewMessageHandler = async (
113113
vscode.window.showInformationMessage(getRouterUnavailableSignInMessage())
114114
}
115115

116+
const getCompletionCheckpointForMessage = (
117+
currentCline: { clineMessages: ClineMessage[] },
118+
checkpointTs?: number,
119+
) => {
120+
if (checkpointTs !== undefined) {
121+
const checkpoint = currentCline.clineMessages.find(
122+
(message) =>
123+
message.ts === checkpointTs &&
124+
message.type === "say" &&
125+
message.say === "checkpoint_saved" &&
126+
typeof message.text === "string",
127+
)
128+
129+
return checkpoint ? { ts: checkpoint.ts, commitHash: checkpoint.text! } : undefined
130+
}
131+
132+
return getCompletionCheckpoint(currentCline.clineMessages)
133+
}
134+
116135
const getCurrentMode = async (): Promise<string> => {
117136
const currentTask = provider.getCurrentTask()
118137

@@ -1291,7 +1310,9 @@ export const webviewMessageHandler = async (
12911310
}
12921311
case "completionCheckpointDiff": {
12931312
const currentCline = provider.getCurrentTask()
1294-
const checkpoint = currentCline ? getCompletionCheckpoint(currentCline.clineMessages) : undefined
1313+
const checkpoint = currentCline
1314+
? getCompletionCheckpointForMessage(currentCline, message.checkpointTs)
1315+
: undefined
12951316

12961317
if (currentCline && checkpoint) {
12971318
await currentCline.checkpointDiff({
@@ -1305,7 +1326,9 @@ export const webviewMessageHandler = async (
13051326
}
13061327
case "completionCheckpointRestore": {
13071328
const currentCline = provider.getCurrentTask()
1308-
const checkpoint = currentCline ? getCompletionCheckpoint(currentCline.clineMessages) : undefined
1329+
const checkpoint = currentCline
1330+
? getCompletionCheckpointForMessage(currentCline, message.checkpointTs)
1331+
: undefined
13091332

13101333
if (checkpoint) {
13111334
await provider.cancelTask()
@@ -1314,6 +1337,7 @@ export const webviewMessageHandler = async (
13141337
await pWaitFor(() => provider.getCurrentTask()?.isInitialized === true, { timeout: 3_000 })
13151338
} catch (error) {
13161339
vscode.window.showErrorMessage(t("common:errors.checkpoint_timeout"))
1340+
return
13171341
}
13181342

13191343
try {

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1676,6 +1676,9 @@ export const ChatRowContent = ({
16761676
</div>
16771677
<div style={{ color: "var(--vscode-charts-green)", paddingTop: 10 }}>
16781678
<Markdown markdown={message.text} partial={message.partial} />
1679+
{!message.partial && enableCheckpoints !== false && completionCheckpoint ? (
1680+
<SeeNewChangesButtons checkpoint={completionCheckpoint} />
1681+
) : null}
16791682
</div>
16801683
</div>
16811684
)

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -137,6 +137,11 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
137137
if (message?.type === "say" && message.say === "completion_result") {
138138
return message.ts
139139
}
140+
141+
// Handle ask-only completion_result rows that carry completion text.
142+
if (message?.type === "ask" && message.ask === "completion_result" && (message.text ?? "") !== "") {
143+
return message.ts
144+
}
140145
}
141146

142147
return undefined

‎webview-ui/src/components/chat/SeeNewChangesButtons.tsx‎

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -11,14 +11,14 @@ export const SeeNewChangesButtons = ({ checkpoint }: { checkpoint: CompletionChe
1111
const [restoringChanges, setRestoringChanges] = useState(false)
1212

1313
const seeNewChangesCallback = useCallback(() => {
14-
vscode.postMessage({ type: "completionCheckpointDiff" })
15-
}, [])
14+
vscode.postMessage({ type: "completionCheckpointDiff", checkpointTs: checkpoint.ts })
15+
}, [checkpoint.ts])
1616

1717
const restoreChangesCallback = useCallback(() => setRestoringChanges(true), [])
1818

1919
const confirmRestoreChangesCallback = useCallback(() => {
20-
vscode.postMessage({ type: "completionCheckpointRestore" })
21-
}, [])
20+
vscode.postMessage({ type: "completionCheckpointRestore", checkpointTs: checkpoint.ts })
21+
}, [checkpoint.ts])
2222

2323
const cancelRestoreChangesCallback = useCallback(() => setRestoringChanges(false), [])
2424

@@ -38,12 +38,17 @@ export const SeeNewChangesButtons = ({ checkpoint }: { checkpoint: CompletionChe
3838
</>
3939
) : (
4040
<>
41-
<VSCodeButton className="w-full mt-2" appearance="secondary" onClick={seeNewChangesCallback}>
41+
<VSCodeButton
42+
className="w-full mt-2"
43+
appearance="secondary"
44+
title={t("chat:seeNewChanges.tooltip")}
45+
onClick={seeNewChangesCallback}>
4246
{t("chat:seeNewChanges.title")}
4347
</VSCodeButton>
4448
<VSCodeButton
4549
className="w-full mt-2"
4650
appearance="secondary"
51+
title={t("chat:restoreChanges.tooltip")}
4752
disabled={!checkpoint.commitHash}
4853
onClick={restoreChangesCallback}>
4954
{t("chat:restoreChanges.title")}

‎webview-ui/src/components/chat/__tests__/ChatView.clear-approval-buttons.spec.tsx‎

Lines changed: 31 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,9 @@ vi.mock("../ChatRow", () => ({
4242
return (
4343
<div data-testid="chat-row">
4444
{JSON.stringify(message)}
45-
{message.type === "say" && message.say === "completion_result" && completionCheckpoint ? (
45+
{((message.type === "say" && message.say === "completion_result") ||
46+
(message.type === "ask" && message.ask === "completion_result" && message.text)) &&
47+
completionCheckpoint ? (
4648
<div data-testid="completion-checkpoint-actions">
4749
<button type="button">chat:seeNewChanges.title</button>
4850
<button type="button">chat:restoreChanges.title</button>
@@ -158,12 +160,12 @@ const autoApprovedCommandAsk = (): ClineMessage[] => [
158160
]
159161

160162
const completionAskWithoutCheckpoint = (): ClineMessage[] => [
161-
{ type: "say", say: "text", ts: 1, text: "Initial task" },
163+
{ type: "say", say: "task", ts: 1, text: "Initial task" },
162164
{ type: "ask", ask: "completion_result", ts: 2, text: "Task complete", partial: false },
163165
]
164166

165167
const completionAskWithCheckpoint = (): ClineMessage[] => [
166-
{ type: "say", say: "text", ts: 1, text: "Initial task" },
168+
{ type: "say", say: "task", ts: 1, text: "Initial task" },
167169
{
168170
type: "say",
169171
say: "checkpoint_saved",
@@ -175,6 +177,18 @@ const completionAskWithCheckpoint = (): ClineMessage[] => [
175177
{ type: "ask", ask: "completion_result", ts: 4, text: "", partial: false },
176178
]
177179

180+
const askOnlyCompletionWithCheckpoint = (): ClineMessage[] => [
181+
{ type: "say", say: "task", ts: 1, text: "Initial task" },
182+
{
183+
type: "say",
184+
say: "checkpoint_saved",
185+
ts: 2,
186+
text: "checkpoint-after-user-prompt",
187+
checkpoint: { suppressMessage: true },
188+
},
189+
{ type: "ask", ask: "completion_result", ts: 3, text: "Task complete", partial: false },
190+
]
191+
178192
describe("ChatView approval button behavior", () => {
179193
beforeEach(() => vi.clearAllMocks())
180194

@@ -227,6 +241,20 @@ describe("ChatView approval button behavior", () => {
227241
})
228242
})
229243

244+
it("shows inline checkpoint actions for ask-only completion rows with text", async () => {
245+
const { queryByText, queryByTestId } = renderChatView()
246+
247+
await act(async () => {
248+
hydrateState(askOnlyCompletionWithCheckpoint())
249+
})
250+
251+
await waitFor(() => {
252+
expect(queryByTestId("completion-checkpoint-actions")).toBeInTheDocument()
253+
expect(queryByText(SEE_NEW_CHANGES_BUTTON_LABEL)).toBeInTheDocument()
254+
expect(queryByText(RESTORE_CHANGES_BUTTON_LABEL)).toBeInTheDocument()
255+
})
256+
})
257+
230258
it("keeps Start New Task for completion results without a latest-prompt checkpoint", async () => {
231259
const { queryByText } = renderChatView()
232260

0 commit comments

Comments
 (0)