From bb0aeddaaee7abb4ca076752167e8af14d52683f Mon Sep 17 00:00:00 2001
From: zzgu <guzhizhou28@gmail.com>
Date: Thu, 2 Jan 2025 13:54:00 -0800
Subject: [PATCH] fix: implement shortcut in client middelware

---
 clients/tabby-agent/src/codeLens.ts          |  8 ++-----
 clients/vscode/src/lsp/CodeLensMiddleware.ts | 24 ++++++++++++++++++--
 2 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/clients/tabby-agent/src/codeLens.ts b/clients/tabby-agent/src/codeLens.ts
index e626bd1da470..3340f0ce81c3 100644
--- a/clients/tabby-agent/src/codeLens.ts
+++ b/clients/tabby-agent/src/codeLens.ts
@@ -11,7 +11,6 @@ import {
 import { ClientCapabilities, ServerCapabilities, CodeLens, CodeLensType, ChangesPreviewLineType } from "./protocol";
 import { TextDocuments } from "./lsp/textDocuments";
 import { TextDocument } from "vscode-languageserver-textdocument";
-import { isBrowser } from "./env";
 
 const codeLensType: CodeLensType = "previewChanges";
 const changesPreviewLineType = {
@@ -117,12 +116,10 @@ export class CodeLensProvider implements Feature {
               },
             });
           } else if (!previewBlockMarkers.includes("x")) {
-            // TODO: read keybinds from LSP client, then send to LSP server to avoid hardcode.
-            const acceptShortcut = isBrowser ? "" : ` (${process.platform === "darwin" ? "cmd+enter" : "ctrl+enter"})`;
             lineCodeLenses.push({
               range: codeLensRange,
               command: {
-                title: `$(check)Accept${acceptShortcut}`,
+                title: `$(check)Accept`,
                 command: "tabby/chat/edit/resolve",
                 arguments: [{ location: codeLensLocation, action: "accept" }],
               },
@@ -132,11 +129,10 @@ export class CodeLensProvider implements Feature {
               },
             });
 
-            const discardShortcut = isBrowser ? "" : ` (esc)`;
             lineCodeLenses.push({
               range: codeLensRange,
               command: {
-                title: `$(remove-close)Discard${discardShortcut}`,
+                title: `$(remove-close)Discard`,
                 command: "tabby/chat/edit/resolve",
                 arguments: [{ location: codeLensLocation, action: "discard" }],
               },
diff --git a/clients/vscode/src/lsp/CodeLensMiddleware.ts b/clients/vscode/src/lsp/CodeLensMiddleware.ts
index 9a9cfd28ef84..00dd803f6ff4 100644
--- a/clients/vscode/src/lsp/CodeLensMiddleware.ts
+++ b/clients/vscode/src/lsp/CodeLensMiddleware.ts
@@ -12,6 +12,7 @@ import {
 import { CodeLensMiddleware as VscodeLspCodeLensMiddleware, ProvideCodeLensesSignature } from "vscode-languageclient";
 import { CodeLens as TabbyCodeLens } from "tabby-agent";
 import { findTextEditor } from "./vscodeWindowUtils";
+import { isBrowser } from "../env";
 
 type CodeLens = VscodeCodeLens & TabbyCodeLens;
 
@@ -74,16 +75,22 @@ export class CodeLensMiddleware implements VscodeLspCodeLensMiddleware {
     if (!editor) {
       return codeLenses;
     }
+
+    if (!codeLenses) {
+      return [];
+    }
+
     this.removeDecorations(editor);
     const result =
       codeLenses
-        ?.map((codeLens) => this.handleCodeLens(codeLens, editor))
-        .filter((codeLens): codeLens is CodeLens => codeLens !== null) ?? [];
+        .map((codeLens) => this.handleCodeLens(codeLens, editor))
+        .filter((codeLens): codeLens is CodeLens => codeLens !== null);
     this.purgeDecorationMap();
     return result;
   }
 
   private handleCodeLens(codeLens: CodeLens, editor: TextEditor): CodeLens | null {
+    this.addShortcut(codeLens);
     if (!codeLens.data || codeLens.data.type !== "previewChanges") {
       return codeLens;
     }
@@ -106,6 +113,19 @@ export class CodeLensMiddleware implements VscodeLspCodeLensMiddleware {
     return null;
   }
 
+  private addShortcut(codeLens: CodeLens) {    
+    if (codeLens.command?.arguments?.[0].action === "accept") {
+      // TODO: read keybinds from LSP client, then send to LSP server to avoid hardcode.
+      const acceptShortcut = isBrowser ? '' : ` (${process.platform === 'darwin' ? 'cmd+enter' : 'ctrl+enter'})`;
+
+      codeLens.command.title += (acceptShortcut);
+    } else if (codeLens.command?.arguments?.[0].action === "discard") {
+      const discardShortcut = isBrowser ? '' : ` (esc)`;
+
+      codeLens.command.title += (discardShortcut);
+    }
+  }
+
   private addDecorationRange(editor: TextEditor, decorationType: TextEditorDecorationType, range: Range) {
     let decorations: Map<TextEditorDecorationType, Range[]> | undefined;
     if (this.decorationMap.has(editor)) {