Let CodeLensProvider using footer's markers

TabbyML · Aug 28, 2024 · 46c120a · 46c120a
1 parent 619cf43
commit 46c120a
Show file tree

Hide file tree

Showing 2 changed files with 67 additions and 33 deletions.
diff --git a/clients/tabby-agent/src/lsp/ChatEditProvider.ts b/clients/tabby-agent/src/lsp/ChatEditProvider.ts
@@ -23,6 +23,7 @@ import * as Diff from "diff";
 import { TabbyAgent } from "../TabbyAgent";
 import { isEmptyRange } from "../utils/range";
 import { isBlank } from "../utils";
+import { getLogger } from "../logger";
 
 export type Edit = {
   id: ChatEditToken;
@@ -187,21 +188,26 @@ export class ChatEditProvider {
     if (!document) {
       return false;
     }
-    const header = document.getText({
-      start: {
-        line: params.location.range.start.line,
-        character: 0,
-      },
-      end: {
-        line: params.location.range.start.line + 1,
-        character: 0,
-      },
-    });
-    const match = /^<<<<<<<.+(<.*>)\[(tabby-[0-9|a-z|A-Z]{6})\]/g.exec(header);
-    const markers = match?.[1];
-    if (!match || !markers) {
-      return false;
+
+    let markers;
+    let line = params.location.range.start.line;
+    for (; line < document.lineCount; line++) {
+      const lineText = document.getText({
+        start: { line, character: 0 },
+        end: { line: line + 1, character: 0 },
+      });
+
+      const match = /^>>>>>>>.+(<.*>)\[(tabby-[0-9|a-z|A-Z]{6})\]/g.exec(lineText);
+      markers = match?.[1];
+      if (markers) {
+        break;
+      }
     }
+
+    if (!markers) {
+      return false
+    }
+
     const previewRange = {
       start: {
         line: params.location.range.start.line,
@@ -257,6 +263,32 @@ export class ChatEditProvider {
     responseCommentTag?: string[],
   ): Promise<void> {
     const applyEdit = async (edit: Edit, isFirst: boolean = false, isLast: boolean = false) => {
+      if (isFirst) {
+        const workspaceEdit: WorkspaceEdit = {
+          changes: {
+            [edit.location.uri]: [
+              {
+                range: edit.editedRange,
+                newText: `<<<<<<< Inline Edit <>[${edit.id}]\n`
+              },
+            ],
+          },
+        };
+
+        await this.applyWorkspaceEdit({
+          edit: workspaceEdit,
+          options: {
+            undoStopBefore: isFirst,
+            undoStopAfter: isLast,
+          },
+        });
+
+        edit.editedRange = {
+          start: { line: edit.editedRange.start.line + 1, character: 0 },
+          end: { line: edit.editedRange.start.line + 1, character: 0 },
+        };
+      }
+
       const editedLines = this.generateChangesPreview(edit);
       const workspaceEdit: WorkspaceEdit = {
         changes: {
@@ -407,7 +439,7 @@ export class ChatEditProvider {
     } else if (edit.state == "completed") {
       stateDescription = "Editing completed";
     }
-    lines.push(`<<<<<<< ${stateDescription} {{markers}}[${edit.id}]`);
+    // lines.push(`<<<<<<< ${stateDescription} {{markers}}[${edit.id}]`);
     markers += "<";
     // comments: split by new line or 80 chars
     const commentLines = edit.comments
@@ -485,7 +517,7 @@ export class ChatEditProvider {
     lines.push(`>>>>>>> ${stateDescription} {{markers}}[${edit.id}]`);
     markers += ">";
     // replace markers
-    lines[0] = lines[0]!.replace("{{markers}}", markers);
+    // lines[0] = lines[0]!.replace("{{markers}}", markers);
     lines[lines.length - 1] = lines[lines.length - 1]!.replace("{{markers}}", markers);
     return lines;
   }

diff --git a/clients/tabby-agent/src/lsp/CodeLensProvider.ts b/clients/tabby-agent/src/lsp/CodeLensProvider.ts
@@ -10,6 +10,7 @@ import {
 import { ServerCapabilities, CodeLens, CodeLensType, ChangesPreviewLineType } from "./protocol";
 import { TextDocuments } from "./TextDocuments";
 import { TextDocument } from "vscode-languageserver-textdocument";
+import { getLogger } from "../logger";
 
 const codeLensType: CodeLensType = "previewChanges";
 const changesPreviewLineType = {
@@ -54,7 +55,7 @@ export class CodeLensProvider {
     const codeLenses: CodeLens[] = [];
     let lineInPreviewBlock = -1;
     let previewBlockMarkers = "";
-    for (let line = 0; line < textDocument.lineCount; line++) {
+    for (let line = textDocument.lineCount - 1; line >= 0; line = line - 1) {
       if (token.isCancellationRequested) {
         return null;
       }
@@ -64,15 +65,29 @@ export class CodeLensProvider {
         start: { line: line, character: 0 },
         end: { line: line, character: text.length - 1 },
       };
+
       const codeLensLocation: Location = { uri: uri, range: codeLensRange };
       const lineCodeLenses: CodeLens[] = [];
       if (lineInPreviewBlock < 0) {
-        const match = /^<<<<<<<.+(<.*>)\[(tabby-[0-9|a-z|A-Z]{6})\]/g.exec(text);
+        const match = /^>>>>>>>.+(<.*>)\[(tabby-[0-9|a-z|A-Z]{6})\]/g.exec(text);
         const markers = match?.[1];
         const editId = match?.[2];
         if (match && markers && editId) {
-          lineInPreviewBlock = 0;
           previewBlockMarkers = markers;
+          lineInPreviewBlock = 0;
+          lineCodeLenses.push({
+            range: codeLensRange,
+            data: {
+              type: codeLensType,
+              line: changesPreviewLineType.footer,
+            },
+          });
+        }
+      } else {
+        const match = /^<<<<<<<.+(<.*>)\[(tabby-[0-9|a-z|A-Z]{6})\]/g.exec(text);
+        const editId = match?.[2];
+        if (match && editId) {
+          lineInPreviewBlock = -1;
 
           lineCodeLenses.push({
             range: codeLensRange,
@@ -98,22 +113,9 @@ export class CodeLensProvider {
               line: changesPreviewLineType.header,
             },
           });
-        }
-      } else {
-        const match = /^>>>>>>>.+(<.*>)\[(tabby-[0-9|a-z|A-Z]{6})\]/g.exec(text);
-        const editId = match?.[2];
-        if (match && editId) {
-          lineInPreviewBlock = -1;
-          lineCodeLenses.push({
-            range: codeLensRange,
-            data: {
-              type: codeLensType,
-              line: changesPreviewLineType.footer,
-            },
-          });
         } else {
           lineInPreviewBlock++;
-          const marker = previewBlockMarkers[lineInPreviewBlock];
+          const marker = previewBlockMarkers[previewBlockMarkers.length - lineInPreviewBlock - 1];
           let codeLens: CodeLens | undefined = undefined;
           switch (marker) {
             case "#":