Skip to content

Commit 07edf21

Browse files
robertpiosikrobertpiosik
committed
Implement temporary checkpoints and promotion logic to enable safe undo/revert functionality during the chat response application and review process
1 parent 5739e81 commit 07edf21

5 files changed

Lines changed: 136 additions & 97 deletions

File tree

‎packages/ui/src/components/editor/panel/Changes/Changes.tsx‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -220,7 +220,7 @@ export const Changes: FC<Props> = ({
220220
is_focused={is_keep_button_focused}
221221
>
222222
<span className="codicon codicon-check" />
223-
{files.length > 1 ? 'Keep Selected' : 'Keep'}
223+
{files.length > 1 ? 'Approve Selected' : 'Approve'}
224224
</Button>
225225
</div>
226226
</div>

‎packages/vscode/src/commands/apply-chat-response-command/apply-chat-response-command.ts‎

Lines changed: 54 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -17,32 +17,30 @@ import {
1717
ModelProvidersManager,
1818
get_tool_config_id
1919
} from '@/services/model-providers-manager'
20+
import {
21+
create_temporary_checkpoint,
22+
delete_checkpoint,
23+
promote_temporary_checkpoint
24+
} from '@/commands/checkpoints-command'
2025
import { apply_git_patch } from './handlers/diff-handler'
2126
import { PROVIDERS } from '@shared/constants/providers'
2227
import { LAST_SELECTED_INTELLIGENT_UPDATE_CONFIG_ID_STATE_KEY } from '@/constants/state-keys'
2328
import { Diff } from './utils/clipboard-parser/extract-diff-patches'
2429
import { ViewProvider } from '@/views/panel/backend/panel-provider'
2530
import { review, code_review_promise_resolve } from './utils/review'
2631
import { dictionary } from '@shared/constants/dictionary'
32+
import { WorkspaceProvider } from '@/context/providers/workspace-provider'
2733

2834
let ongoing_review_cleanup_promise: Promise<void> | null = null
2935

3036
const check_if_all_files_new = async (
3137
files: ClipboardFile[]
3238
): Promise<boolean> => {
33-
if (
34-
!vscode.workspace.workspaceFolders ||
35-
vscode.workspace.workspaceFolders.length == 0
36-
) {
37-
return false
38-
}
39-
4039
const workspace_map = new Map<string, string>()
41-
vscode.workspace.workspaceFolders.forEach((folder) => {
40+
vscode.workspace.workspaceFolders!.forEach((folder) => {
4241
workspace_map.set(folder.name, folder.uri.fsPath)
4342
})
44-
45-
const default_workspace = vscode.workspace.workspaceFolders[0].uri.fsPath
43+
const default_workspace = vscode.workspace.workspaceFolders![0].uri.fsPath
4644

4745
for (const file of files) {
4846
let workspace_root = default_workspace
@@ -333,7 +331,8 @@ const handle_code_review_and_cleanup = async (params: {
333331

334332
export const apply_chat_response_command = (
335333
context: vscode.ExtensionContext,
336-
view_provider: ViewProvider
334+
view_provider: ViewProvider,
335+
workspace_provider: WorkspaceProvider
337336
) => {
338337
const update_undo_and_apply_button_state = (
339338
states: OriginalFileState[] | null,
@@ -382,20 +381,11 @@ export const apply_chat_response_command = (
382381
position: { line: number; character: number }
383382
}
384383
}) => {
385-
if (args?.raw_instructions) {
386-
Logger.info({
387-
function_name: 'apply_chat_response_command',
388-
message: 'Received raw instructions',
389-
data: { raw_instructions: args.raw_instructions }
390-
})
391-
}
392-
393-
if (args?.edit_format) {
394-
Logger.info({
395-
function_name: 'apply_chat_response_command',
396-
message: 'Received edit format',
397-
data: { edit_format: args.edit_format }
398-
})
384+
if (!vscode.workspace.workspaceFolders?.length) {
385+
vscode.window.showErrorMessage(
386+
dictionary.error_message.NO_WORKSPACE_FOLDER_OPEN
387+
)
388+
return
399389
}
400390

401391
if (code_review_promise_resolve) {
@@ -416,6 +406,10 @@ export const apply_chat_response_command = (
416406
}
417407
}
418408

409+
const temp_checkpoint = await create_temporary_checkpoint(
410+
workspace_provider
411+
)
412+
419413
type ReviewData = {
420414
original_states: OriginalFileState[]
421415
chat_response: string
@@ -452,21 +446,12 @@ export const apply_chat_response_command = (
452446
clipboard_content.code_completion
453447
) {
454448
const completion = clipboard_content.code_completion
455-
if (
456-
!vscode.workspace.workspaceFolders ||
457-
vscode.workspace.workspaceFolders.length == 0
458-
) {
459-
vscode.window.showErrorMessage(
460-
dictionary.error_message.NO_WORKSPACE_FOLDER_OPEN
461-
)
462-
return null
463-
}
464449
const workspace_map = new Map<string, string>()
465-
vscode.workspace.workspaceFolders.forEach((folder) => {
450+
vscode.workspace.workspaceFolders!.forEach((folder) => {
466451
workspace_map.set(folder.name, folder.uri.fsPath)
467452
})
468453
const default_workspace =
469-
vscode.workspace.workspaceFolders[0].uri.fsPath
454+
vscode.workspace.workspaceFolders![0].uri.fsPath
470455
let workspace_root = default_workspace
471456
if (
472457
completion.workspace_name &&
@@ -536,13 +521,6 @@ export const apply_chat_response_command = (
536521
}
537522

538523
if (clipboard_content.type == 'patches' && clipboard_content.patches) {
539-
if (!vscode.workspace.workspaceFolders?.length) {
540-
vscode.window.showErrorMessage(
541-
dictionary.error_message.NO_WORKSPACE_FOLDER_OPEN
542-
)
543-
return null
544-
}
545-
546524
const rename_map = new Map<string, string>()
547525
clipboard_content.patches.forEach((patch) => {
548526
if (patch.new_file_path && patch.file_path) {
@@ -562,12 +540,12 @@ export const apply_chat_response_command = (
562540
}
563541

564542
const workspace_map = new Map<string, string>()
565-
vscode.workspace.workspaceFolders.forEach((folder) => {
543+
vscode.workspace.workspaceFolders!.forEach((folder) => {
566544
workspace_map.set(folder.name, folder.uri.fsPath)
567545
})
568546

569547
const default_workspace =
570-
vscode.workspace.workspaceFolders[0].uri.fsPath
548+
vscode.workspace.workspaceFolders![0].uri.fsPath
571549

572550
let success_count = 0
573551
let failure_count = 0
@@ -900,15 +878,37 @@ export const apply_chat_response_command = (
900878
}
901879
})()
902880

903-
if (review_data) {
904-
await handle_code_review_and_cleanup({
905-
original_states: review_data.original_states,
906-
chat_response: review_data.chat_response,
907-
view_provider,
908-
update_undo_and_apply_button_state,
909-
original_editor_state: args?.original_editor_state,
910-
raw_instructions: args?.raw_instructions
911-
})
881+
let checkpoint_promoted = false
882+
883+
try {
884+
if (review_data) {
885+
const changes_accepted = await handle_code_review_and_cleanup({
886+
original_states: review_data.original_states,
887+
chat_response: review_data.chat_response,
888+
view_provider,
889+
update_undo_and_apply_button_state,
890+
original_editor_state: args?.original_editor_state,
891+
raw_instructions: args?.raw_instructions
892+
})
893+
894+
if (changes_accepted) {
895+
await promote_temporary_checkpoint({
896+
context,
897+
temp_checkpoint,
898+
title: 'Before approve changes',
899+
description: args?.raw_instructions
900+
})
901+
checkpoint_promoted = true
902+
}
903+
}
904+
} finally {
905+
if (!checkpoint_promoted) {
906+
await delete_checkpoint({
907+
context,
908+
checkpoint_to_delete: temp_checkpoint,
909+
options: { skip_undo_prompt: true }
910+
})
911+
}
912912
}
913913
}
914914
)

packages/vscode/src/commands/checkpoints-command.ts renamed to packages/vscode/src/commands/checkpoints-command/checkpoints-command.ts

Lines changed: 79 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -5,15 +5,15 @@ import * as fs from 'fs/promises'
55
import {
66
CHECKPOINTS_STATE_KEY,
77
TEMPORARY_CHECKPOINT_TIMESTAMP_STATE_KEY
8-
} from '../constants/state-keys'
9-
import { WorkspaceProvider } from '../context/providers/workspace-provider'
10-
import { should_ignore_file } from '../context/utils/should-ignore-file'
8+
} from '../../constants/state-keys'
9+
import { WorkspaceProvider } from '../../context/providers/workspace-provider'
10+
import { should_ignore_file } from '../../context/utils/should-ignore-file'
1111
import dayjs from 'dayjs'
1212
import relativeTime from 'dayjs/plugin/relativeTime'
1313

1414
dayjs.extend(relativeTime)
1515

16-
interface Checkpoint {
16+
export interface Checkpoint {
1717
timestamp: number
1818
title: string
1919
description?: string
@@ -224,13 +224,13 @@ const create_checkpoint = async (
224224
}
225225
}
226226

227-
const create_temporary_checkpoint = async (
227+
export const create_temporary_checkpoint = async (
228228
workspace_provider: WorkspaceProvider
229-
): Promise<Checkpoint> =>
230-
vscode.window.withProgress(
229+
): Promise<Checkpoint> => {
230+
return vscode.window.withProgress(
231231
{
232232
location: vscode.ProgressLocation.Notification,
233-
title: 'Warming up...',
233+
title: 'Preparing checkpoint...',
234234
cancellable: false
235235
},
236236
async () => {
@@ -250,6 +250,27 @@ const create_temporary_checkpoint = async (
250250
return new_checkpoint
251251
}
252252
)
253+
}
254+
255+
export const promote_temporary_checkpoint = async (params: {
256+
context: vscode.ExtensionContext
257+
temp_checkpoint: Checkpoint
258+
title: string
259+
description?: string
260+
}) => {
261+
const checkpoints =
262+
params.context.workspaceState.get<Checkpoint[]>(
263+
CHECKPOINTS_STATE_KEY,
264+
[]
265+
) ?? []
266+
checkpoints.push({
267+
...params.temp_checkpoint,
268+
is_temporary: false,
269+
title: params.title,
270+
description: params.description
271+
})
272+
await params.context.workspaceState.update(CHECKPOINTS_STATE_KEY, checkpoints)
273+
}
253274

254275
const sync_directory = async (params: {
255276
source_dir: vscode.Uri
@@ -529,7 +550,7 @@ const restore_checkpoint = async (params: {
529550
await delete_checkpoint({
530551
context: params.context,
531552
checkpoint_to_delete: temp_checkpoint,
532-
options: { is_reverting: true }
553+
options: { skip_undo_prompt: true }
533554
})
534555
await params.context.workspaceState.update(
535556
TEMPORARY_CHECKPOINT_TIMESTAMP_STATE_KEY,
@@ -547,7 +568,7 @@ const restore_checkpoint = async (params: {
547568
await delete_checkpoint({
548569
context: params.context,
549570
checkpoint_to_delete: temp_checkpoint,
550-
options: { is_reverting: true }
571+
options: { skip_undo_prompt: true }
551572
})
552573
await params.context.workspaceState.update(
553574
TEMPORARY_CHECKPOINT_TIMESTAMP_STATE_KEY,
@@ -557,34 +578,12 @@ const restore_checkpoint = async (params: {
557578
}
558579
}
559580

560-
const delete_checkpoint = async (params: {
581+
export const delete_checkpoint = async (params: {
561582
context: vscode.ExtensionContext
562583
checkpoint_to_delete: Checkpoint
563-
options?: { is_reverting?: boolean }
584+
options?: { skip_undo_prompt?: boolean }
564585
}) => {
565-
if (!params.options?.is_reverting) {
566-
const confirmation = await vscode.window.showWarningMessage(
567-
`Are you sure you want to delete the checkpoint from ${dayjs(
568-
params.checkpoint_to_delete.timestamp
569-
).fromNow()}? This action cannot be undone.`,
570-
{ modal: true },
571-
'Delete'
572-
)
573-
574-
if (confirmation != 'Delete') return
575-
}
576-
577-
try {
578-
const checkpoint_path = get_checkpoint_path(
579-
params.checkpoint_to_delete.timestamp
580-
)
581-
await vscode.workspace.fs.delete(vscode.Uri.file(checkpoint_path), {
582-
recursive: true
583-
})
584-
} catch (error) {
585-
console.warn(`Could not delete checkpoint file: ${error}`)
586-
}
587-
586+
// First, remove checkpoint from state. This makes the deletion immediate from user's perspective.
588587
const checkpoints =
589588
params.context.workspaceState.get<Checkpoint[]>(
590589
CHECKPOINTS_STATE_KEY,
@@ -598,12 +597,51 @@ const delete_checkpoint = async (params: {
598597
updated_checkpoints
599598
)
600599

601-
if (!params.options?.is_reverting) {
602-
vscode.window.showInformationMessage(
603-
`Checkpoint from ${dayjs(
600+
const actually_delete_files = async () => {
601+
try {
602+
const checkpoint_path = get_checkpoint_path(
604603
params.checkpoint_to_delete.timestamp
605-
).fromNow()} deleted successfully.`
606-
)
604+
)
605+
await vscode.workspace.fs.delete(vscode.Uri.file(checkpoint_path), {
606+
recursive: true
607+
})
608+
} catch (error) {
609+
console.warn(`Could not delete checkpoint file: ${error}`)
610+
}
611+
}
612+
613+
if (!params.options?.skip_undo_prompt) {
614+
vscode.window
615+
.showInformationMessage(
616+
`Checkpoint from ${dayjs(
617+
params.checkpoint_to_delete.timestamp
618+
).fromNow()} deleted successfully.`,
619+
'Revert'
620+
)
621+
.then(async (action) => {
622+
if (action == 'Revert') {
623+
// User chose to revert. Add the checkpoint back to the state.
624+
const current_checkpoints =
625+
params.context.workspaceState.get<Checkpoint[]>(
626+
CHECKPOINTS_STATE_KEY,
627+
[]
628+
) ?? []
629+
current_checkpoints.push(params.checkpoint_to_delete)
630+
await params.context.workspaceState.update(
631+
CHECKPOINTS_STATE_KEY,
632+
current_checkpoints
633+
)
634+
vscode.window.showInformationMessage('Checkpoint deletion reverted.')
635+
} else {
636+
// User dismissed the notification or it timed out.
637+
// Proceed with deleting the checkpoint files.
638+
await actually_delete_files()
639+
}
640+
})
641+
} else {
642+
// This is a silent deletion (e.g. temporary checkpoint cleanup), no need
643+
// to ask for revert.
644+
await actually_delete_files()
607645
}
608646
}
609647

@@ -789,7 +827,7 @@ export const checkpoints_command = (
789827
await delete_checkpoint({
790828
context,
791829
checkpoint_to_delete: temp_checkpoint,
792-
options: { is_reverting: true }
830+
options: { skip_undo_prompt: true }
793831
})
794832
await context.workspaceState.update(
795833
TEMPORARY_CHECKPOINT_TIMESTAMP_STATE_KEY,

‎packages/vscode/src/commands/checkpoints-command/index.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
export * from './checkpoints-command'

‎packages/vscode/src/extension.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,7 @@ export async function activate(context: vscode.ExtensionContext) {
6969
}
7070
),
7171
reference_in_chat_command(view_provider, workspace_provider),
72-
apply_chat_response_command(context, view_provider),
72+
apply_chat_response_command(context, view_provider, workspace_provider),
7373
apply_code_block_to_active_editor_command(context),
7474
undo_command(
7575
context,

0 commit comments

Comments
 (0)