1 files changed, 22 insertions, 14 deletions

diff --git a/extension/src/continueIdeClient.ts b/extension/src/continueIdeClient.ts
index 08a0b74d..8ab3e075 100644
--- a/extension/src/continueIdeClient.ts
+++ b/extension/src/continueIdeClient.ts
@@ -1,5 +1,8 @@
 // import { ShowSuggestionRequest } from "../schema/ShowSuggestionRequest";
-import { showSuggestion, SuggestionRanges } from "./suggestions";
+import {
+  showSuggestion as showSuggestionInEditor,
+  SuggestionRanges,
+} from "./suggestions";
 import { openEditorAndRevealRange, getRightViewColumn } from "./util/vscode";
 import { FileEdit } from "../schema/FileEdit";
 import { RangeInFile } from "../schema/RangeInFile";
@@ -114,7 +117,10 @@ class IdeProtocolClient {
         break;
       case "setFileOpen":
         this.openFile(data.filepath);
-        // TODO: Close file
+        // TODO: Close file if False
+        break;
+      case "showSuggestion":
+        this.showSuggestion(data.edit);
         break;
       case "openGUI":
       case "connected":
@@ -138,6 +144,7 @@ class IdeProtocolClient {
   // ------------------------------------ //
   // On message handlers
 
+  private _lastDecorationType: vscode.TextEditorDecorationType | null = null;
   async highlightCode(rangeInFile: RangeInFile, color: string) {
     const range = new vscode.Range(
       rangeInFile.range.start.line,
@@ -157,24 +164,25 @@ class IdeProtocolClient {
       });
       editor.setDecorations(decorationType, [range]);
 
-      // Listen for changes to cursor position and then remove the decoration (but keep for at least 2 seconds)
-      const allowRemoveHighlight = () => {
-        const cursorDisposable = vscode.window.onDidChangeTextEditorSelection(
-          (event) => {
-            if (event.textEditor.document.uri.fsPath === rangeInFile.filepath) {
-              cursorDisposable.dispose();
-              editor.setDecorations(decorationType, []);
-            }
+      const cursorDisposable = vscode.window.onDidChangeTextEditorSelection(
+        (event) => {
+          if (event.textEditor.document.uri.fsPath === rangeInFile.filepath) {
+            cursorDisposable.dispose();
+            editor.setDecorations(decorationType, []);
           }
-        );
-      };
-      setTimeout(allowRemoveHighlight, 2000);
+        }
+      );
+
+      if (this._lastDecorationType) {
+        editor.setDecorations(this._lastDecorationType, []);
+      }
+      this._lastDecorationType = decorationType;
     }
   }
 
   showSuggestion(edit: FileEdit) {
     // showSuggestion already exists
-    showSuggestion(
+    showSuggestionInEditor(
       edit.filepath,
       new vscode.Range(
         edit.range.start.line,