side-by-side diff editor

6 files changed, 221 insertions, 67 deletions

diff --git a/continuedev/src/continuedev/steps/core/core.py b/continuedev/src/continuedev/steps/core/core.py
index b9f0da35..b215b317 100644
--- a/continuedev/src/continuedev/steps/core/core.py
+++ b/continuedev/src/continuedev/steps/core/core.py
@@ -295,6 +295,17 @@ class DefaultModelEditCodeStep(Step):
         prompt = self.compile_prompt(file_prefix, contents, file_suffix, sdk)
         full_file_contents_lines = full_file_contents.split("\n")
 
+        async def sendDiffUpdate(lines: List[str], sdk: ContinueSDK):
+            nonlocal full_file_contents_lines, rif
+
+            completion = "\n".join(lines)
+
+            full_prefix_lines = full_file_contents_lines[:rif.range.start.line]
+            full_suffix_lines = full_file_contents_lines[rif.range.end.line + 1:]
+            new_file_contents = "\n".join(
+                full_prefix_lines) + "\n" + completion + "\n" + "\n".join(full_suffix_lines)
+            await sdk.ide.showDiff(rif.filepath, new_file_contents)
+
         # Important state variables
         # -------------------------
         original_lines = [] if rif.contents == "" else rif.contents.split("\n")
@@ -435,16 +446,16 @@ class DefaultModelEditCodeStep(Step):
                 chunk_lines.pop()  # because this will be an empty string
             else:
                 unfinished_line = chunk_lines.pop()
-            lines.extend(chunk_lines)
+            lines.extend(map(lambda l: common_whitespace + l, chunk_lines))
+
+            if True:
+                await sendDiffUpdate(lines, sdk)
 
             # Deal with newly accumulated lines
             for line in chunk_lines:
                 # Trailing whitespace doesn't matter
                 line = line.rstrip()
 
-                # Add the common whitespace that was removed before prompting
-                line = common_whitespace + line
-
                 # Lines that should signify the end of generation
                 if self.is_end_line(line):
                     break
@@ -463,7 +474,9 @@ class DefaultModelEditCodeStep(Step):
                     break
 
                 # If none of the above, insert the line!
-                await handle_generated_line(line)
+                if False:
+                    await handle_generated_line(line)
+
                 completion_lines_covered += 1
                 current_line_in_file += 1
 
@@ -475,34 +488,37 @@ class DefaultModelEditCodeStep(Step):
             completion_lines_covered += 1
             current_line_in_file += 1
 
-        # If the current block isn't empty, add that suggestion
-        if len(current_block_lines) > 0:
-            # We have a chance to back-track here for blank lines that are repeats of the end of the original
-            # Don't want to have the same ending in both the original and the generated, can just leave it there
-            num_to_remove = 0
-            for i in range(-1, -len(current_block_lines) - 1, -1):
-                if len(original_lines_below_previous_blocks) == 0:
-                    break
-                if current_block_lines[i] == original_lines_below_previous_blocks[-1]:
-                    num_to_remove += 1
-                    original_lines_below_previous_blocks.pop()
-                else:
-                    break
-            current_block_lines = current_block_lines[:-
-                                                      num_to_remove] if num_to_remove > 0 else current_block_lines
-
-            # It's also possible that some lines match at the beginning of the block
-            # while len(current_block_lines) > 0 and len(original_lines_below_previous_blocks) > 0 and current_block_lines[0] == original_lines_below_previous_blocks[0]:
-            #     current_block_lines.pop(0)
-            #     original_lines_below_previous_blocks.pop(0)
-            #     current_block_start += 1
-
-            await sdk.ide.showSuggestion(FileEdit(
-                filepath=rif.filepath,
-                range=Range.from_shorthand(
-                    current_block_start, 0, current_block_start + len(original_lines_below_previous_blocks), 0),
-                replacement="\n".join(current_block_lines)
-            ))
+        await sendDiffUpdate(lines, sdk)
+
+        if False:
+            # If the current block isn't empty, add that suggestion
+            if len(current_block_lines) > 0:
+                # We have a chance to back-track here for blank lines that are repeats of the end of the original
+                # Don't want to have the same ending in both the original and the generated, can just leave it there
+                num_to_remove = 0
+                for i in range(-1, -len(current_block_lines) - 1, -1):
+                    if len(original_lines_below_previous_blocks) == 0:
+                        break
+                    if current_block_lines[i] == original_lines_below_previous_blocks[-1]:
+                        num_to_remove += 1
+                        original_lines_below_previous_blocks.pop()
+                    else:
+                        break
+                current_block_lines = current_block_lines[:-
+                                                          num_to_remove] if num_to_remove > 0 else current_block_lines
+
+                # It's also possible that some lines match at the beginning of the block
+                # while len(current_block_lines) > 0 and len(original_lines_below_previous_blocks) > 0 and current_block_lines[0] == original_lines_below_previous_blocks[0]:
+                #     current_block_lines.pop(0)
+                #     original_lines_below_previous_blocks.pop(0)
+                #     current_block_start += 1
+
+                await sdk.ide.showSuggestion(FileEdit(
+                    filepath=rif.filepath,
+                    range=Range.from_shorthand(
+                        current_block_start, 0, current_block_start + len(original_lines_below_previous_blocks), 0),
+                    replacement="\n".join(current_block_lines)
+                ))
 
         # Record the completion
         completion = "\n".join(lines)
diff --git a/extension/package.json b/extension/package.json
index 87dd7ba6..461f5721 100644
--- a/extension/package.json
+++ b/extension/package.json
@@ -87,6 +87,16 @@
         "title": "Reject Suggestion"
       },
       {
+        "command": "continue.acceptDiff",
+        "category": "Continue",
+        "title": "Accept Diff"
+      },
+      {
+        "command": "continue.rejectDiff",
+        "category": "Continue",
+        "title": "Reject Diff"
+      },
+      {
         "command": "continue.acceptAllSuggestions",
         "category": "Continue",
         "title": "Accept All Suggestions"
diff --git a/extension/src/commands.ts b/extension/src/commands.ts
index 8072353b..4414a171 100644
--- a/extension/src/commands.ts
+++ b/extension/src/commands.ts
@@ -12,6 +12,8 @@ import {
   acceptAllSuggestionsCommand,
   rejectAllSuggestionsCommand,
 } from "./suggestions";
+
+import { acceptDiffCommand, rejectDiffCommand } from "./diffs";
 import * as bridge from "./bridge";
 import { debugPanelWebview } from "./debugPanel";
 import { sendTelemetryEvent, TelemetryEvent } from "./telemetry";
@@ -51,6 +53,8 @@ const commandsMap: { [command: string]: (...args: any) => any } = {
   "continue.suggestionUp": suggestionUpCommand,
   "continue.acceptSuggestion": acceptSuggestionCommand,
   "continue.rejectSuggestion": rejectSuggestionCommand,
+  "continue.acceptDiff": acceptDiffCommand,
+  "continue.rejectDiff": rejectDiffCommand,
   "continue.acceptAllSuggestions": acceptAllSuggestionsCommand,
   "continue.rejectAllSuggestions": rejectAllSuggestionsCommand,
   "continue.focusContinueInput": async () => {
diff --git a/extension/src/continueIdeClient.ts b/extension/src/continueIdeClient.ts
index b9969858..90547edc 100644
--- a/extension/src/continueIdeClient.ts
+++ b/extension/src/continueIdeClient.ts
@@ -15,6 +15,10 @@ import {
 import { FileEditWithFullContents } from "../schema/FileEditWithFullContents";
 import fs = require("fs");
 import { WebsocketMessenger } from "./util/messenger";
+import * as path from "path";
+import * as os from "os";
+import { diffManager } from "./diffs";
+
 class IdeProtocolClient {
   private messenger: WebsocketMessenger | null = null;
   private readonly context: vscode.ExtensionContext;
@@ -239,40 +243,8 @@ class IdeProtocolClient {
     );
   }
 
-  contentProvider: vscode.Disposable | null = null;
-
   showDiff(filepath: string, replacement: string) {
-    const myProvider = new (class
-      implements vscode.TextDocumentContentProvider
-    {
-      onDidChangeEmitter = new vscode.EventEmitter<vscode.Uri>();
-      onDidChange = this.onDidChangeEmitter.event;
-      provideTextDocumentContent = (uri: vscode.Uri) => {
-        return replacement;
-      };
-    })();
-    this.contentProvider = vscode.workspace.registerTextDocumentContentProvider(
-      "continueDiff",
-      myProvider
-    );
-
-    // Call the event fire
-    const diffFilename = `continueDiff://${filepath}`;
-    myProvider.onDidChangeEmitter.fire(vscode.Uri.parse(diffFilename));
-
-    const leftUri = vscode.Uri.file(filepath);
-    const rightUri = vscode.Uri.parse(diffFilename);
-    const title = "Continue Diff";
-    vscode.commands
-      .executeCommand("vscode.diff", leftUri, rightUri, title)
-      .then(
-        () => {
-          console.log("Diff view opened successfully");
-        },
-        (error) => {
-          console.error("Error opening diff view:", error);
-        }
-      );
+    diffManager.writeDiff(filepath, replacement);
   }
 
   openFile(filepath: string) {
diff --git a/extension/src/diffs.ts b/extension/src/diffs.ts
new file mode 100644
index 00000000..4bd072cf
--- /dev/null
+++ b/extension/src/diffs.ts
@@ -0,0 +1,114 @@
+import * as os from "os";
+import * as path from "path";
+import * as fs from "fs";
+import * as vscode from "vscode";
+
+interface DiffInfo {
+  originalFilepath: string;
+  newFilepath: string;
+  editor?: vscode.TextEditor;
+}
+
+export const DIFF_DIRECTORY = path.join(os.homedir(), ".continue", "diffs");
+
+class DiffManager {
+  // Create a temporary file in the global .continue directory which displays the updated version
+  // Doing this because virtual files are read-only
+  private diffs: Map<string, DiffInfo> = new Map();
+
+  constructor() {
+    // Make sure the diff directory exists
+    if (!fs.existsSync(DIFF_DIRECTORY)) {
+      fs.mkdirSync(DIFF_DIRECTORY, {
+        recursive: true,
+      });
+    }
+  }
+
+  private escapeFilepath(filepath: string): string {
+    return filepath.replace(/\\/g, "_").replace(/\//g, "_");
+  }
+
+  private openDiffEditor(
+    originalFilepath: string,
+    newFilepath: string,
+    newContent: string
+  ): vscode.TextEditor {
+    const rightUri = vscode.Uri.parse(newFilepath);
+    const leftUri = vscode.Uri.file(originalFilepath);
+    const title = "Continue Diff";
+    vscode.commands.executeCommand("vscode.diff", leftUri, rightUri, title);
+
+    const editor = vscode.window.activeTextEditor;
+    if (!editor) {
+      throw new Error("No active text editor found for Continue Diff");
+    }
+    return editor;
+  }
+
+  writeDiff(originalFilepath: string, newContent: string): string {
+    // Create or update existing diff
+    const newFilepath = path.join(
+      DIFF_DIRECTORY,
+      this.escapeFilepath(originalFilepath)
+    );
+    fs.writeFileSync(newFilepath, newContent);
+
+    // Open the diff editor if this is a new diff
+    if (!this.diffs.has(originalFilepath)) {
+      const diffInfo: DiffInfo = {
+        originalFilepath,
+        newFilepath,
+      };
+      diffInfo.editor = this.openDiffEditor(
+        originalFilepath,
+        newFilepath,
+        newContent
+      );
+      this.diffs.set(originalFilepath, diffInfo);
+    }
+    return newFilepath;
+  }
+
+  cleanUpDiff(diffInfo: DiffInfo) {
+    // Close the editor, remove the record, delete the file
+    if (diffInfo.editor) {
+      vscode.window.showTextDocument(diffInfo.editor.document);
+      vscode.commands.executeCommand("workbench.action.closeActiveEditor");
+    }
+    this.diffs.delete(diffInfo.originalFilepath);
+    fs.unlinkSync(diffInfo.newFilepath);
+  }
+
+  acceptDiff(originalFilepath: string) {
+    // Get the diff info, copy new file to original, then delete from record and close the corresponding editor
+    const diffInfo = this.diffs.get(originalFilepath);
+    if (!diffInfo) {
+      return;
+    }
+    fs.writeFileSync(
+      diffInfo.originalFilepath,
+      fs.readFileSync(diffInfo.newFilepath)
+    );
+    this.cleanUpDiff(diffInfo);
+  }
+
+  rejectDiff(originalFilepath: string) {
+    const diffInfo = this.diffs.get(originalFilepath);
+    if (!diffInfo) {
+      return;
+    }
+
+    this.cleanUpDiff(diffInfo);
+  }
+}
+
+export const diffManager = new DiffManager();
+
+export async function acceptDiffCommand(originalFilepath: string) {
+  diffManager.acceptDiff(originalFilepath);
+}
+
+export async function rejectDiffCommand(originalFilepath: string) {
+  diffManager.rejectDiff(originalFilepath);
+}
diff --git a/extension/src/lang-server/codeLens.ts b/extension/src/lang-server/codeLens.ts
index 3bd4f153..08435a3b 100644
--- a/extension/src/lang-server/codeLens.ts
+++ b/extension/src/lang-server/codeLens.ts
@@ -1,6 +1,8 @@
 import * as vscode from "vscode";
 import { editorToSuggestions, editorSuggestionsLocked } from "../suggestions";
-
+import * as path from "path";
+import * as os from "os";
+import { DIFF_DIRECTORY } from "../diffs";
 class SuggestionsCodeLensProvider implements vscode.CodeLensProvider {
   public provideCodeLenses(
     document: vscode.TextDocument,
@@ -60,15 +62,51 @@ class SuggestionsCodeLensProvider implements vscode.CodeLensProvider {
   }
 }
 
+class DiffViewerCodeLensProvider implements vscode.CodeLensProvider {
+  public provideCodeLenses(
+    document: vscode.TextDocument,
+    token: vscode.CancellationToken
+  ): vscode.CodeLens[] | Thenable<vscode.CodeLens[]> {
+    if (path.dirname(document.uri.fsPath) !== DIFF_DIRECTORY) {
+      return [];
+    } else {
+      const codeLenses: vscode.CodeLens[] = [];
+      const range = new vscode.Range(0, 0, 0, 0);
+      codeLenses.push(
+        new vscode.CodeLens(range, {
+          title: "Accept ✅",
+          command: "continue.acceptDiff",
+          arguments: [document.uri.fsPath],
+        }),
+        new vscode.CodeLens(range, {
+          title: "Reject ❌",
+          command: "continue.rejectDiff",
+          arguments: [document.uri.fsPath],
+        })
+      );
+      return codeLenses;
+    }
+  }
+}
+
+let diffsCodeLensDisposable: vscode.Disposable | undefined = undefined;
 let suggestionsCodeLensDisposable: vscode.Disposable | undefined = undefined;
 
 export function registerAllCodeLensProviders(context: vscode.ExtensionContext) {
   if (suggestionsCodeLensDisposable) {
     suggestionsCodeLensDisposable.dispose();
   }
+  if (diffsCodeLensDisposable) {
+    diffsCodeLensDisposable.dispose();
+  }
   suggestionsCodeLensDisposable = vscode.languages.registerCodeLensProvider(
     "*",
     new SuggestionsCodeLensProvider()
   );
+  diffsCodeLensDisposable = vscode.languages.registerCodeLensProvider(
+    "*",
+    new DiffViewerCodeLensProvider()
+  );
   context.subscriptions.push(suggestionsCodeLensDisposable);
+  context.subscriptions.push(diffsCodeLensDisposable);
 }