side-by-side diff editor

4 files changed, 162 insertions, 34 deletions

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);
 }