feat: add checkpoints + undo/redo for edit

core/index.d.ts

@@ -200,11 +200,16 @@ export interface IContextProvider {
   loadSubmenuItems(args: LoadSubmenuItemsArgs): Promise<ContextSubmenuItem[]>;
 }
 
+export interface Checkpoint {
+  [filepath: string]: string;
+}
+
 export interface PersistedSessionInfo {
   history: ChatHistory;
   title: string;
   workspaceDirectory: string;
   sessionId: string;
+  checkpoints?: Checkpoint[];
 }
 
 export interface SessionInfo {

core/protocol/ideWebview.ts

@@ -27,6 +27,7 @@ export type ToIdeFromWebviewProtocol = ToIdeFromWebviewOrCoreProtocol & {
     },
     void,
   ];
+  overwriteFile: [{ filepath: string; prevFileContent: string | null }, void];
   showTutorial: [undefined, void];
   showFile: [{ filepath: string }, void];
   openConfigJson: [undefined, void];
@@ -40,8 +41,8 @@ export type ToIdeFromWebviewProtocol = ToIdeFromWebviewOrCoreProtocol & {
   "jetbrains/isOSREnabled": [undefined, void];
   "vscode/openMoveRightMarkdown": [undefined, void];
   setGitHubAuthToken: [{ token: string }, void];
-  acceptDiff: [{ filepath: string }, void];
-  rejectDiff: [{ filepath: string }, void];
+  acceptDiff: [{ filepath: string; streamId?: string }, void];
+  rejectDiff: [{ filepath: string; streamId?: string }, void];
   "edit/sendPrompt": [
     { prompt: MessageContent; range: RangeInFileWithContents },
     void,
@@ -64,11 +65,17 @@ export type EditStatus =
   | "accepting:full-diff"
   | "done";
 
+export type ApplyStateStatus =
+  | "streaming" // Changes are being applied to the file
+  | "done" // All changes have been applied, awaiting user to accept/reject
+  | "closed"; // All changes have been applied. Note that for new files, we immediately set the status to "closed"
+
 export interface ApplyState {
   streamId: string;
-  status?: "streaming" | "done" | "closed";
+  status?: ApplyStateStatus;
   numDiffs?: number;
   filepath?: string;
+  fileContent?: string;
 }
 
 export type ToWebviewFromIdeProtocol = ToWebviewFromIdeOrCoreProtocol & {

extensions/vscode/src/commands.ts

@@ -271,7 +271,10 @@ const commandsMap: (
   }
 
   return {
-    "continue.acceptDiff": async (newFilepath?: string | vscode.Uri) => {
+    "continue.acceptDiff": async (
+      newFilepath?: string | vscode.Uri,
+      streamId?: string,
+    ) => {
       captureCommandTelemetry("acceptDiff");
 
       let fullPath = newFilepath;
@@ -290,8 +293,23 @@ const commandsMap: (
       await sidebar.webviewProtocol.request("setEditStatus", {
         status: "done",
       });
+
+      if (streamId && fullPath) {
+        const fileContent = await ide.readFile(fullPath);
+
+        await sidebar.webviewProtocol.request("updateApplyState", {
+          fileContent,
+          filepath: fullPath,
+          streamId: streamId,
+          status: "closed",
+          numDiffs: 0,
+        });
+      }
     },
-    "continue.rejectDiff": async (newFilepath?: string | vscode.Uri) => {
+    "continue.rejectDiff": async (
+      newFilepath?: string | vscode.Uri,
+      streamId?: string,
+    ) => {
       captureCommandTelemetry("rejectDiff");
 
       let fullPath = newFilepath;
@@ -309,6 +327,18 @@ const commandsMap: (
       await sidebar.webviewProtocol.request("setEditStatus", {
         status: "done",
       });
+
+      if (streamId && fullPath) {
+        const fileContent = await ide.readFile(fullPath);
+
+        await sidebar.webviewProtocol.request("updateApplyState", {
+          fileContent,
+          filepath: fullPath,
+          streamId: streamId,
+          status: "closed",
+          numDiffs: 0,
+        });
+      }
     },
     "continue.acceptVerticalDiffBlock": (filepath?: string, index?: number) => {
       captureCommandTelemetry("acceptVerticalDiffBlock");

extensions/vscode/src/diff/vertical/handler.ts

@@ -18,6 +18,7 @@ export interface VerticalDiffHandlerOptions {
   onStatusUpdate: (
     status?: ApplyState["status"],
     numDiffs?: ApplyState["numDiffs"],
+    fileContent?: ApplyState["fileContent"],
   ) => void;
 }
 
@@ -271,6 +272,7 @@ export class VerticalDiffHandler implements vscode.Disposable {
     this.options.onStatusUpdate(
       "closed",
       this.editorToVerticalDiffCodeLens.get(this.filepath)?.length ?? 0,
+      this.editor.document.getText(),
     );
 
     this.cancelled = true;
@@ -362,6 +364,7 @@ export class VerticalDiffHandler implements vscode.Disposable {
       this.options.onStatusUpdate(
         "done",
         this.editorToVerticalDiffCodeLens.get(this.filepath)?.length ?? 0,
+        this.editor.document.getText(),
       );
 
       // Reject on user typing
@@ -416,7 +419,11 @@ export class VerticalDiffHandler implements vscode.Disposable {
       this.editorToVerticalDiffCodeLens.get(this.filepath)?.length ?? 0;
 
     const status = numDiffs === 0 ? "closed" : undefined;
-    this.options.onStatusUpdate(status, numDiffs);
+    this.options.onStatusUpdate(
+      status,
+      numDiffs,
+      this.editor.document.getText(),
+    );
   }
 
   private shiftCodeLensObjects(startLine: number, offset: number) {

extensions/vscode/src/diff/vertical/manager.ts

@@ -217,11 +217,13 @@
       endLine,
       {
         instant,
-        onStatusUpdate: (status, numDiffs) =>
+        onStatusUpdate: (status, numDiffs, fileContent) =>
           this.webviewProtocol.request("updateApplyState", {
             streamId,
             status,
             numDiffs,
+            fileContent,
+            filepath,
           }),
       },
     );
@@ -298,8 +300,8 @@
         // Previous diff was a quickEdit
         // Check if user has highlighted a range
         let rangeBool =
           startLine != endLine ||
           editor.selection.start.character != editor.selection.end.character;
 
         // Check if the range is different from the previous range
         let newRangeBool =
@@ -336,12 +338,14 @@
       endLine,
       {
         input,
-        onStatusUpdate: (status, numDiffs) =>
+        onStatusUpdate: (status, numDiffs, fileContent) =>
           streamId &&
           this.webviewProtocol.request("updateApplyState", {
             streamId,
             status,
             numDiffs,
+            fileContent,
+            filepath,
           }),
       },
     );

extensions/vscode/src/extension/VsCodeMessenger.ts

@@ -152,15 +152,35 @@ export class VsCodeMessenger {
       );
     });
 
-    webviewProtocol.on("acceptDiff", async ({ data: { filepath } }) => {
-      await vscode.commands.executeCommand("continue.acceptDiff", filepath);
-    });
+    webviewProtocol.on(
+      "acceptDiff",
+      async ({ data: { filepath, streamId } }) => {
+        await vscode.commands.executeCommand(
+          "continue.acceptDiff",
+          filepath,
+          streamId,
+        );
+      },
+    );
 
-    webviewProtocol.on("rejectDiff", async ({ data: { filepath } }) => {
-      await vscode.commands.executeCommand("continue.rejectDiff", filepath);
-    });
+    webviewProtocol.on(
+      "rejectDiff",
+      async ({ data: { filepath, streamId } }) => {
+        await vscode.commands.executeCommand(
+          "continue.rejectDiff",
+          filepath,
+          streamId,
+        );
+      },
+    );
 
     this.onWebview("applyToFile", async ({ data }) => {
+      webviewProtocol.request("updateApplyState", {
+        streamId: data.streamId,
+        status: "streaming",
+        fileContent: data.text,
+      });
+
       let filepath = data.filepath;
 
       // If there is a filepath, verify it exists and then open the file
@@ -173,18 +193,10 @@ export class VsCodeMessenger {
 
         const fileExists = await this.ide.fileExists(fullPath);
 
-        // If it's a new file, no need to apply, just write directly
+        // If there is no existing file at the path, create it
         if (!fileExists) {
-          await this.ide.writeFile(fullPath, data.text);
+          await this.ide.writeFile(fullPath, "");
           await this.ide.openFile(fullPath);
-
-          await webviewProtocol.request("updateApplyState", {
-            streamId: data.streamId,
-            status: "done",
-            numDiffs: 0,
-          });
-
-          return;
         }
 
         await this.ide.openFile(fullPath);
@@ -203,11 +215,14 @@ export class VsCodeMessenger {
         editor.edit((builder) =>
           builder.insert(new vscode.Position(0, 0), data.text),
         );
+
         await webviewProtocol.request("updateApplyState", {
           streamId: data.streamId,
-          status: "done",
+          status: "closed",
           numDiffs: 0,
+          fileContent: data.text,
         });
+
         return;
       }
 
@@ -295,6 +310,37 @@ export class VsCodeMessenger {
     this.onWebview("openUrl", (msg) => {
       vscode.env.openExternal(vscode.Uri.parse(msg.data));
     });
+
+    this.onWebview(
+      "overwriteFile",
+      async ({ data: { prevFileContent, filepath } }) => {
+        if (prevFileContent === null) {
+          // TODO: Delete the file
+          return;
+        }
+
+        await this.ide.openFile(filepath);
+
+        // Get active text editor
+        const editor = vscode.window.activeTextEditor;
+
+        if (!editor) {
+          vscode.window.showErrorMessage("No active editor to apply edits to");
+          return;
+        }
+
+        editor.edit((builder) =>
+          builder.replace(
+            new vscode.Range(
+              editor.document.positionAt(0),
+              editor.document.positionAt(editor.document.getText().length),
+            ),
+            prevFileContent,
+          ),
+        );
+      },
+    );
+
     this.onWebview("insertAtCursor", async (msg) => {
       const editor = vscode.window.activeTextEditor;
       if (editor === undefined || !editor.selection) {

gui/src/components/Layout.tsx

@@ -12,7 +12,11 @@ import {
   setEditStatus,
   startEditMode,
 } from "../redux/slices/editModeState";
-import { setShowDialog, updateApplyState } from "../redux/slices/uiStateSlice";
+import { setShowDialog } from "../redux/slices/uiStateSlice";
+import {
+  updateApplyState,
+  updateCurCheckpoint,
+} from "../redux/slices/stateSlice";
 import { RootState } from "../redux/store";
 import { getFontSize, isMetaEquivalentKeyPressed } from "../util";
 import { getLocalStorage, setLocalStorage } from "../util/localStorage";
@@ -158,6 +162,12 @@ const Layout = () => {
   useWebviewListener(
     "updateApplyState",
     async (state) => {
+      dispatch(
+        updateCurCheckpoint({
+          filepath: state.filepath,
+          content: state.fileContent,
+        }),
+      );
       dispatch(updateApplyState(state));
     },
     [],

gui/src/components/StepContainer/AcceptRejectAllButtons.tsx

@@ -0,0 +1,46 @@
+import { useContext } from "react";
+import { IdeMessengerContext } from "../../context/IdeMessenger";
+import { CheckIcon, XMarkIcon } from "@heroicons/react/24/outline";
+import { ApplyState } from "core/protocol/ideWebview";
+
+export interface AcceptRejectAllButtonsProps {
+  pendingApplyStates: ApplyState[];
+}
+
+export default function AcceptRejectAllButtons({
+  pendingApplyStates,
+}: AcceptRejectAllButtonsProps) {
+  const ideMessenger = useContext(IdeMessengerContext);
+
+  async function handleAcceptOrReject(status: "acceptDiff" | "rejectDiff") {
+    for (const { filepath, streamId } of pendingApplyStates) {
+      ideMessenger.post(status, {
+        filepath,
+        streamId,
+      });
+    }
+  }
+
+  return (
+    <div className="flex justify-center gap-2 border-b border-gray-200/25 p-1 px-3">
+      <button
+        className="flex cursor-pointer items-center border-none bg-transparent px-2 py-1 text-xs text-gray-300 opacity-80 hover:opacity-100 hover:brightness-125"
+        onClick={() => handleAcceptOrReject("rejectDiff")}
+      >
+        <XMarkIcon className="mr-1 h-4 w-4 text-red-600" />
+        <span className="sm:hidden">Reject</span>
+        <span className="max-sm:hidden md:hidden">Reject all</span>
+        <span className="max-md:hidden">Reject all changes</span>
+      </button>
+      <button
+        className="flex cursor-pointer items-center border-none bg-transparent px-2 py-1 text-xs text-gray-300 opacity-80 hover:opacity-100 hover:brightness-125"
+        onClick={() => handleAcceptOrReject("acceptDiff")}
+      >
+        <CheckIcon className="mr-1 h-4 w-4 text-green-600" />
+        <span className="sm:hidden">Accept</span>
+        <span className="max-sm:hidden md:hidden">Accept all</span>
+        <span className="max-md:hidden">Accept all changes</span>
+      </button>
+    </div>
+  );
+}