diff editor infer filepath, codelens in middle

1 files changed, 6 insertions, 2 deletions

diff --git a/extension/src/lang-server/codeLens.ts b/extension/src/lang-server/codeLens.ts
index 778b98dc..5800a00e 100644
--- a/extension/src/lang-server/codeLens.ts
+++ b/extension/src/lang-server/codeLens.ts
@@ -2,7 +2,7 @@ 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";
 class SuggestionsCodeLensProvider implements vscode.CodeLensProvider {
   public provideCodeLenses(
     document: vscode.TextDocument,
@@ -53,7 +53,11 @@ class DiffViewerCodeLensProvider implements vscode.CodeLensProvider {
   ): 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 ✅ (⌘⇧↩)",