1 files changed, 55 insertions, 1 deletions

diff --git a/extension/src/lang-server/codeLens.ts b/extension/src/lang-server/codeLens.ts
index 3bd4f153..381a0084 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,67 @@ 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) {
+      const codeLenses: vscode.CodeLens[] = [];
+      const range = new vscode.Range(0, 0, 1, 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;
+    } 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
+          );
+        }
+      });
+    }
+  }
+}
+
+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);
 }