1 files changed, 123 insertions, 14 deletions

diff --git a/extension/src/diffs.ts b/extension/src/diffs.ts
index 28089fc6..1130a06a 100644
--- a/extension/src/diffs.ts
+++ b/extension/src/diffs.ts
@@ -3,21 +3,30 @@ import * as path from "path";
 import * as fs from "fs";
 import * as vscode from "vscode";
 import { extensionContext, ideProtocolClient } from "./activation/activate";
+import { getMetaKeyLabel } from "./util/util";
+import { devDataPath } from "./activation/environmentSetup";
 
 interface DiffInfo {
   originalFilepath: string;
   newFilepath: string;
   editor?: vscode.TextEditor;
   step_index: number;
+  range: vscode.Range;
 }
 
-export const DIFF_DIRECTORY = path.join(os.homedir(), ".continue", "diffs");
+export const DIFF_DIRECTORY = path
+  .join(os.homedir(), ".continue", "diffs")
+  .replace(/^C:/, "c:");
 
 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();
 
+  diffAtNewFilepath(newFilepath: string): DiffInfo | undefined {
+    return this.diffs.get(newFilepath);
+  }
+
   private setupDirectory() {
     // Make sure the diff directory exists
     if (!fs.existsSync(DIFF_DIRECTORY)) {
@@ -29,12 +38,25 @@ class DiffManager {
 
   constructor() {
     this.setupDirectory();
+
+    // Listen for file closes, and if it's a diff file, clean up
+    vscode.workspace.onDidCloseTextDocument((document) => {
+      const newFilepath = document.uri.fsPath;
+      const diffInfo = this.diffs.get(newFilepath);
+      if (diffInfo) {
+        this.cleanUpDiff(diffInfo, false);
+      }
+    });
   }
 
   private escapeFilepath(filepath: string): string {
     return filepath.replace(/\\/g, "_").replace(/\//g, "_");
   }
 
+  private getNewFilepath(originalFilepath: string): string {
+    return path.join(DIFF_DIRECTORY, this.escapeFilepath(originalFilepath));
+  }
+
   private openDiffEditor(
     originalFilepath: string,
     newFilepath: string
@@ -47,7 +69,7 @@ class DiffManager {
       return undefined;
     }
 
-    const rightUri = vscode.Uri.parse(newFilepath);
+    const rightUri = vscode.Uri.file(newFilepath);
     const leftUri = vscode.Uri.file(originalFilepath);
     const title = "Continue Diff";
     console.log(
@@ -77,7 +99,7 @@ class DiffManager {
     ) {
       vscode.window
         .showInformationMessage(
-          "Accept (⌘⇧↩) or reject (⌘⇧⌫) at the top of the file.",
+          `Accept (${getMetaKeyLabel()}⇧↩) or reject (${getMetaKeyLabel()}⇧⌫) at the top of the file.`,
           "Got it",
           "Don't show again"
         )
@@ -95,6 +117,17 @@ class DiffManager {
     return editor;
   }
 
+  private _findFirstDifferentLine(contentA: string, contentB: string): number {
+    const linesA = contentA.split("\n");
+    const linesB = contentB.split("\n");
+    for (let i = 0; i < linesA.length && i < linesB.length; i++) {
+      if (linesA[i] !== linesB[i]) {
+        return i;
+      }
+    }
+    return 0;
+  }
+
   writeDiff(
     originalFilepath: string,
     newContent: string,
@@ -103,18 +136,20 @@ class DiffManager {
     this.setupDirectory();
 
     // Create or update existing diff
-    const newFilepath = path.join(
-      DIFF_DIRECTORY,
-      this.escapeFilepath(originalFilepath)
-    );
+    const newFilepath = this.getNewFilepath(originalFilepath);
     fs.writeFileSync(newFilepath, newContent);
 
     // Open the diff editor if this is a new diff
     if (!this.diffs.has(newFilepath)) {
+      // Figure out the first line that is different
+      const oldContent = ideProtocolClient.readFile(originalFilepath);
+      const line = this._findFirstDifferentLine(oldContent, newContent);
+
       const diffInfo: DiffInfo = {
         originalFilepath,
         newFilepath,
         step_index,
+        range: new vscode.Range(line, 0, line + 1, 0),
       };
       this.diffs.set(newFilepath, diffInfo);
     }
@@ -126,12 +161,17 @@ class DiffManager {
       this.diffs.set(newFilepath, diffInfo);
     }
 
+    vscode.commands.executeCommand(
+      "workbench.action.files.revert",
+      vscode.Uri.file(newFilepath)
+    );
+
     return newFilepath;
   }
 
-  cleanUpDiff(diffInfo: DiffInfo) {
+  cleanUpDiff(diffInfo: DiffInfo, hideEditor: boolean = true) {
     // Close the editor, remove the record, delete the file
-    if (diffInfo.editor) {
+    if (hideEditor && diffInfo.editor) {
       vscode.window.showTextDocument(diffInfo.editor.document);
       vscode.commands.executeCommand("workbench.action.closeActiveEditor");
     }
@@ -139,10 +179,38 @@ class DiffManager {
     fs.unlinkSync(diffInfo.newFilepath);
   }
 
+  private inferNewFilepath() {
+    const activeEditorPath =
+      vscode.window.activeTextEditor?.document.uri.fsPath;
+    if (activeEditorPath && path.dirname(activeEditorPath) === DIFF_DIRECTORY) {
+      return activeEditorPath;
+    }
+    const visibleEditors = vscode.window.visibleTextEditors.map(
+      (editor) => editor.document.uri.fsPath
+    );
+    for (const editorPath of visibleEditors) {
+      if (path.dirname(editorPath) === DIFF_DIRECTORY) {
+        for (const otherEditorPath of visibleEditors) {
+          if (
+            path.dirname(otherEditorPath) !== DIFF_DIRECTORY &&
+            this.getNewFilepath(otherEditorPath) === editorPath
+          ) {
+            return editorPath;
+          }
+        }
+      }
+    }
+
+    if (this.diffs.size === 1) {
+      return Array.from(this.diffs.keys())[0];
+    }
+    return undefined;
+  }
+
   acceptDiff(newFilepath?: string) {
-    // If no newFilepath is provided and there is only one in the dictionary, use that
-    if (!newFilepath && this.diffs.size === 1) {
-      newFilepath = Array.from(this.diffs.keys())[0];
+    // When coming from a keyboard shortcut, we have to infer the newFilepath from visible text editors
+    if (!newFilepath) {
+      newFilepath = this.inferNewFilepath();
     }
     if (!newFilepath) {
       console.log("No newFilepath provided to accept the diff");
@@ -166,12 +234,14 @@ class DiffManager {
         );
         this.cleanUpDiff(diffInfo);
       });
+
+    recordAcceptReject(true, diffInfo);
   }
 
   rejectDiff(newFilepath?: string) {
     // If no newFilepath is provided and there is only one in the dictionary, use that
-    if (!newFilepath && this.diffs.size === 1) {
-      newFilepath = Array.from(this.diffs.keys())[0];
+    if (!newFilepath) {
+      newFilepath = this.inferNewFilepath();
     }
     if (!newFilepath) {
       console.log(
@@ -195,11 +265,50 @@ class DiffManager {
       .then(() => {
         this.cleanUpDiff(diffInfo);
       });
+
+    recordAcceptReject(false, diffInfo);
   }
 }
 
 export const diffManager = new DiffManager();
 
+function recordAcceptReject(accepted: boolean, diffInfo: DiffInfo) {
+  const collectOn = vscode.workspace
+    .getConfiguration("continue")
+    .get<boolean>("dataSwitch");
+
+  if (collectOn) {
+    const devDataDir = devDataPath();
+    const suggestionsPath = path.join(devDataDir, "suggestions.json");
+
+    // Initialize suggestions list
+    let suggestions = [];
+
+    // Check if suggestions.json exists
+    if (fs.existsSync(suggestionsPath)) {
+      const rawData = fs.readFileSync(suggestionsPath, "utf-8");
+      suggestions = JSON.parse(rawData);
+    }
+
+    // Add the new suggestion to the list
+    suggestions.push({
+      accepted,
+      timestamp: Date.now(),
+      suggestion: diffInfo.originalFilepath,
+    });
+
+    // Send the suggestion to the server
+    // ideProtocolClient.sendAcceptRejectSuggestion(accepted);
+
+    // Write the updated suggestions back to the file
+    fs.writeFileSync(
+      suggestionsPath,
+      JSON.stringify(suggestions, null, 4),
+      "utf-8"
+    );
+  }
+}
+
 export async function acceptDiffCommand(newFilepath?: string) {
   diffManager.acceptDiff(newFilepath);
   ideProtocolClient.sendAcceptRejectDiff(true);