2 files changed, 65 insertions, 39 deletions

diff --git a/extension/src/lang-server/codeActions.ts b/extension/src/lang-server/codeActions.ts
new file mode 100644
index 00000000..07cf5f4e
--- /dev/null
+++ b/extension/src/lang-server/codeActions.ts
@@ -0,0 +1,53 @@
+import * as vscode from "vscode";
+
+class ContinueQuickFixProvider implements vscode.CodeActionProvider {
+  public static readonly providedCodeActionKinds = [
+    vscode.CodeActionKind.QuickFix,
+  ];
+
+  provideCodeActions(
+    document: vscode.TextDocument,
+    range: vscode.Range | vscode.Selection,
+    context: vscode.CodeActionContext,
+    token: vscode.CancellationToken
+  ): vscode.ProviderResult<(vscode.Command | vscode.CodeAction)[]> {
+    if (context.diagnostics.length === 0) {
+      return [];
+    }
+
+    const createQuickFix = (edit: boolean) => {
+      const diagnostic = context.diagnostics[0];
+      const quickFix = new vscode.CodeAction(
+        edit ? "Fix with Continue" : "Ask Continue",
+        vscode.CodeActionKind.QuickFix
+      );
+      quickFix.isPreferred = false;
+      const surroundingRange = new vscode.Range(
+        range.start.translate(-3, 0),
+        range.end.translate(3, 0)
+      );
+      quickFix.command = {
+        command: "continue.quickFix",
+        title: "Continue Quick Fix",
+        arguments: [
+          diagnostic.message,
+          document.getText(surroundingRange),
+          edit,
+        ],
+      };
+      return quickFix;
+    };
+    return [createQuickFix(true), createQuickFix(false)];
+  }
+}
+
+export default function registerQuickFixProvider() {
+  // In your extension's activate function:
+  vscode.languages.registerCodeActionsProvider(
+    { language: "*" },
+    new ContinueQuickFixProvider(),
+    {
+      providedCodeActionKinds: ContinueQuickFixProvider.providedCodeActionKinds,
+    }
+  );
+}
diff --git a/extension/src/lang-server/codeLens.ts b/extension/src/lang-server/codeLens.ts
index 79126eaa..ba80e557 100644
--- a/extension/src/lang-server/codeLens.ts
+++ b/extension/src/lang-server/codeLens.ts
@@ -2,11 +2,12 @@ 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";
+import { DIFF_DIRECTORY, diffManager } from "../diffs";
+import { getMetaKeyLabel } from "../util/util";
 class SuggestionsCodeLensProvider implements vscode.CodeLensProvider {
   public provideCodeLenses(
     document: vscode.TextDocument,
-    token: vscode.CancellationToken
+    _: vscode.CancellationToken
   ): vscode.CodeLens[] | Thenable<vscode.CodeLens[]> {
     const suggestions = editorToSuggestions.get(document.uri.toString());
     if (!suggestions) {
@@ -35,7 +36,7 @@ class SuggestionsCodeLensProvider implements vscode.CodeLensProvider {
       if (codeLenses.length === 2) {
         codeLenses.push(
           new vscode.CodeLens(range, {
-            title: "(⌘⇧↩/⌘⇧⌫ to accept/reject all)",
+            title: `(${getMetaKeyLabel()}⇧↩/${getMetaKeyLabel()}⇧⌫ to accept/reject all)`,
             command: "",
           })
         );
@@ -44,40 +45,28 @@ class SuggestionsCodeLensProvider implements vscode.CodeLensProvider {
 
     return codeLenses;
   }
-
-  onDidChangeCodeLenses?: vscode.Event<void> | undefined;
-
-  constructor(emitter?: vscode.EventEmitter<void>) {
-    if (emitter) {
-      this.onDidChangeCodeLenses = emitter.event;
-      this.onDidChangeCodeLenses(() => {
-        if (vscode.window.activeTextEditor) {
-          this.provideCodeLenses(
-            vscode.window.activeTextEditor.document,
-            new vscode.CancellationTokenSource().token
-          );
-        }
-      });
-    }
-  }
 }
 
 class DiffViewerCodeLensProvider implements vscode.CodeLensProvider {
   public provideCodeLenses(
     document: vscode.TextDocument,
-    token: vscode.CancellationToken
+    _: vscode.CancellationToken
   ): vscode.CodeLens[] | Thenable<vscode.CodeLens[]> {
     if (path.dirname(document.uri.fsPath) === DIFF_DIRECTORY) {
       const codeLenses: vscode.CodeLens[] = [];
-      const range = new vscode.Range(0, 0, 1, 0);
+      let range = new vscode.Range(0, 0, 1, 0);
+      const diffInfo = diffManager.diffAtNewFilepath(document.uri.fsPath);
+      if (diffInfo) {
+        range = diffInfo.range;
+      }
       codeLenses.push(
         new vscode.CodeLens(range, {
-          title: "Accept ✅ (⌘⇧↩)",
+          title: `Accept All ✅ (${getMetaKeyLabel()}⇧↩)`,
           command: "continue.acceptDiff",
           arguments: [document.uri.fsPath],
         }),
         new vscode.CodeLens(range, {
-          title: "Reject ❌ (⌘⇧⌫)",
+          title: `Reject All ❌ (${getMetaKeyLabel()}⇧⌫)`,
           command: "continue.rejectDiff",
           arguments: [document.uri.fsPath],
         })
@@ -87,22 +76,6 @@ class DiffViewerCodeLensProvider implements vscode.CodeLensProvider {
       return [];
     }
   }
-
-  onDidChangeCodeLenses?: vscode.Event<void> | undefined;
-
-  constructor(emitter?: vscode.EventEmitter<void>) {
-    if (emitter) {
-      this.onDidChangeCodeLenses = emitter.event;
-      this.onDidChangeCodeLenses(() => {
-        if (vscode.window.activeTextEditor) {
-          this.provideCodeLenses(
-            vscode.window.activeTextEditor.document,
-            new vscode.CancellationTokenSource().token
-          );
-        }
-      });
-    }
-  }
 }
 
 let diffsCodeLensDisposable: vscode.Disposable | undefined = undefined;