setting to show codelens in diff editor

2 files changed, 26 insertions, 4 deletions

diff --git a/extension/src/diffs.ts b/extension/src/diffs.ts
index b70c3b59..178b1a9d 100644
--- a/extension/src/diffs.ts
+++ b/extension/src/diffs.ts
@@ -47,6 +47,12 @@ class DiffManager {
     if (!editor) {
       throw new Error("No active text editor found for Continue Diff");
     }
+
+    // Change the vscode setting to allow codeLens in diff editor
+    vscode.workspace
+      .getConfiguration("diffEditor", editor.document.uri)
+      .update("codeLens", true, vscode.ConfigurationTarget.Global);
+
     return editor;
   }
 
diff --git a/extension/src/lang-server/codeLens.ts b/extension/src/lang-server/codeLens.ts
index 08435a3b..381a0084 100644
--- a/extension/src/lang-server/codeLens.ts
+++ b/extension/src/lang-server/codeLens.ts
@@ -67,11 +67,9 @@ class DiffViewerCodeLensProvider implements vscode.CodeLensProvider {
     document: vscode.TextDocument,
     token: vscode.CancellationToken
   ): vscode.CodeLens[] | Thenable<vscode.CodeLens[]> {
-    if (path.dirname(document.uri.fsPath) !== DIFF_DIRECTORY) {
-      return [];
-    } else {
+    if (path.dirname(document.uri.fsPath) === DIFF_DIRECTORY) {
       const codeLenses: vscode.CodeLens[] = [];
-      const range = new vscode.Range(0, 0, 0, 0);
+      const range = new vscode.Range(0, 0, 1, 0);
       codeLenses.push(
         new vscode.CodeLens(range, {
           title: "Accept ✅",
@@ -85,6 +83,24 @@ class DiffViewerCodeLensProvider implements vscode.CodeLensProvider {
         })
       );
       return codeLenses;
+    } else {
+      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
+          );
+        }
+      });
     }
   }
 }