the fix. and others.

10 files changed, 118 insertions, 93 deletions

diff --git a/extension/.gitignore b/extension/.gitignore
index 6d1b35bf..43b10889 100644
--- a/extension/.gitignore
+++ b/extension/.gitignore
@@ -3,4 +3,5 @@ node_modules/
 out/
 .vscode-test/
 data/
-src/client
\ No newline at end of file
+src/client
+.continue
diff --git a/extension/package-lock.json b/extension/package-lock.json
index 26e1a631..7565f480 100644
--- a/extension/package-lock.json
+++ b/extension/package-lock.json
@@ -1,12 +1,12 @@
 {
   "name": "continue",
-  "version": "0.0.91",
+  "version": "0.0.97",
   "lockfileVersion": 2,
   "requires": true,
   "packages": {
     "": {
       "name": "continue",
-      "version": "0.0.91",
+      "version": "0.0.97",
       "license": "Apache-2.0",
       "dependencies": {
         "@electron/rebuild": "^3.2.10",
diff --git a/extension/package.json b/extension/package.json
index 314e540f..08f5f081 100644
--- a/extension/package.json
+++ b/extension/package.json
@@ -14,7 +14,7 @@
   "displayName": "Continue",
   "pricing": "Free",
   "description": "The open-source coding autopilot",
-  "version": "0.0.91",
+  "version": "0.0.97",
   "publisher": "Continue",
   "engines": {
     "vscode": "^1.74.0"
@@ -210,7 +210,7 @@
     "lint": "eslint src --ext ts",
     "test": "node ./out/test/runTest.js",
     "package": "cp ./config/prod_config.json ./config/config.json && mkdir -p ./build && vsce package --out ./build && cp ./config/dev_config.json ./config/config.json",
-    "full-package": "cd ../continuedev && poetry build && cp ./dist/continuedev-0.1.1-py3-none-any.whl ../extension/scripts/continuedev-0.1.1-py3-none-any.whl && cd ../extension && npm install && npm run typegen && npm run clientgen && cd react-app && npm install && npm run build && cd .. && npm run package",
+    "full-package": "cd ../continuedev && poetry build && cp ./dist/continuedev-0.1.2-py3-none-any.whl ../extension/scripts/continuedev-0.1.2-py3-none-any.whl && cd ../extension && npm install && npm run typegen && npm run clientgen && cd react-app && npm install && npm run build && cd .. && npm run package",
     "install-extension": "code --install-extension ./build/continue-0.0.8.vsix",
     "uninstall": "code --uninstall-extension .continue",
     "reinstall": "rm -rf ./build && npm run package && npm run uninstall && npm run install-extension"
diff --git a/extension/scripts/requirements.txt b/extension/scripts/requirements.txt
index e6ba1e0a..c51c9d73 100644
--- a/extension/scripts/requirements.txt
+++ b/extension/scripts/requirements.txt
@@ -3,4 +3,4 @@
 # typer==0.7.0
 # pydantic
 # pytest
-./continuedev-0.1.1-py3-none-any.whl
\ No newline at end of file
+./continuedev-0.1.2-py3-none-any.whl
\ No newline at end of file
diff --git a/extension/src/activation/activate.ts b/extension/src/activation/activate.ts
index cd8f0cf3..18650561 100644
--- a/extension/src/activation/activate.ts
+++ b/extension/src/activation/activate.ts
@@ -20,6 +20,21 @@ export async function activateExtension(
 ) {
   extensionContext = context;
 
+  await new Promise((resolve, reject) => {
+    vscode.window.withProgress(
+      {
+        location: vscode.ProgressLocation.Notification,
+        title:
+          "Starting Continue Server... (it may take a minute to download Python packages)",
+        cancellable: false,
+      },
+      async (progress, token) => {
+        await startContinuePythonServer();
+        resolve(null);
+      }
+    );
+  });
+
   sendTelemetryEvent(TelemetryEvent.ExtensionActivated);
   registerAllCodeLensProviders(context);
   registerAllCommands(context);
@@ -33,7 +48,7 @@ export async function activateExtension(
 
   // Setup the left panel
   (async () => {
-    const sessionIdPromise = ideProtocolClient.getSessionId();
+    const sessionIdPromise = await ideProtocolClient.getSessionId();
     const provider = new ContinueGUIWebviewViewProvider(sessionIdPromise);
 
     context.subscriptions.push(
@@ -46,21 +61,6 @@ export async function activateExtension(
       )
     );
   })();
-
-  await new Promise((resolve, reject) => {
-    vscode.window.withProgress(
-      {
-        location: vscode.ProgressLocation.Notification,
-        title:
-          "Starting Continue Server... (it may take a minute to download Python packages)",
-        cancellable: false,
-      },
-      async (progress, token) => {
-        await startContinuePythonServer();
-        resolve(null);
-      }
-    );
-  });
   // All opened terminals should be replaced by our own terminal
   // vscode.window.onDidOpenTerminal((terminal) => {});
 
diff --git a/extension/src/continueIdeClient.ts b/extension/src/continueIdeClient.ts
index b179cbf3..21104abe 100644
--- a/extension/src/continueIdeClient.ts
+++ b/extension/src/continueIdeClient.ts
@@ -32,8 +32,8 @@ class IdeProtocolClient {
     messenger.onClose(() => {
       this.messenger = null;
     });
-    messenger.onMessage((messageType, data) => {
-      this.handleMessage(messageType, data);
+    messenger.onMessage((messageType, data, messenger) => {
+      this.handleMessage(messageType, data, messenger);
     });
 
     // Setup listeners for any file changes in open editors
@@ -67,41 +67,45 @@ class IdeProtocolClient {
     // });
   }
 
-  async handleMessage(messageType: string, data: any) {
+  async handleMessage(
+    messageType: string,
+    data: any,
+    messenger: WebsocketMessenger
+  ) {
     switch (messageType) {
       case "highlightedCode":
-        this.messenger?.send("highlightedCode", {
+        messenger.send("highlightedCode", {
           highlightedCode: this.getHighlightedCode(),
         });
         break;
       case "workspaceDirectory":
-        this.messenger?.send("workspaceDirectory", {
+        messenger.send("workspaceDirectory", {
           workspaceDirectory: this.getWorkspaceDirectory(),
         });
         break;
       case "uniqueId":
-        this.messenger?.send("uniqueId", {
+        messenger.send("uniqueId", {
           uniqueId: this.getUniqueId(),
         });
         break;
       case "getUserSecret":
-        this.messenger?.send("getUserSecret", {
+        messenger.send("getUserSecret", {
           value: await this.getUserSecret(data.key),
         });
         break;
       case "openFiles":
-        this.messenger?.send("openFiles", {
+        messenger.send("openFiles", {
           openFiles: this.getOpenFiles(),
         });
         break;
       case "readFile":
-        this.messenger?.send("readFile", {
+        messenger.send("readFile", {
           contents: this.readFile(data.filepath),
         });
         break;
       case "editFile":
         const fileEdit = await this.editFile(data.edit);
-        this.messenger?.send("editFile", {
+        messenger.send("editFile", {
           fileEdit,
         });
         break;
@@ -109,7 +113,7 @@ class IdeProtocolClient {
         this.highlightCode(data.rangeInFile, data.color);
         break;
       case "runCommand":
-        this.messenger?.send("runCommand", {
+        messenger.send("runCommand", {
           output: await this.runCommand(data.command),
         });
         break;
@@ -249,6 +253,8 @@ class IdeProtocolClient {
         ) {
           clearInterval(interval);
           resolve(null);
+        } else {
+          console.log("Websocket not yet open, trying again...");
         }
       }, 1000);
     });
diff --git a/extension/src/debugPanel.ts b/extension/src/debugPanel.ts
index b176eee7..b88c86f3 100644
--- a/extension/src/debugPanel.ts
+++ b/extension/src/debugPanel.ts
@@ -132,7 +132,7 @@ let streamManager = new StreamManager();
 export let debugPanelWebview: vscode.Webview | undefined;
 export function setupDebugPanel(
   panel: vscode.WebviewPanel | vscode.WebviewView,
-  sessionIdPromise: Promise<string>
+  sessionIdPromise: Promise<string> | string
 ): string {
   debugPanelWebview = panel.webview;
   panel.onDidDispose(() => {
@@ -230,7 +230,12 @@ export function setupDebugPanel(
   panel.webview.onDidReceiveMessage(async (data) => {
     switch (data.type) {
       case "onLoad": {
-        const sessionId = await sessionIdPromise;
+        let sessionId: string;
+        if (typeof sessionIdPromise === "string") {
+          sessionId = sessionIdPromise;
+        } else {
+          sessionId = await sessionIdPromise;
+        }
         panel.webview.postMessage({
           type: "onLoad",
           vscMachineId: vscode.env.machineId,
@@ -350,9 +355,9 @@ export class ContinueGUIWebviewViewProvider
   implements vscode.WebviewViewProvider
 {
   public static readonly viewType = "continue.continueGUIView";
-  private readonly sessionIdPromise: Promise<string>;
+  private readonly sessionIdPromise: Promise<string> | string;
 
-  constructor(sessionIdPromise: Promise<string>) {
+  constructor(sessionIdPromise: Promise<string> | string) {
     this.sessionIdPromise = sessionIdPromise;
   }
 
diff --git a/extension/src/lang-server/codeLens.ts b/extension/src/lang-server/codeLens.ts
index 21448e31..3bd4f153 100644
--- a/extension/src/lang-server/codeLens.ts
+++ b/extension/src/lang-server/codeLens.ts
@@ -60,18 +60,15 @@ class SuggestionsCodeLensProvider implements vscode.CodeLensProvider {
   }
 }
 
-const allCodeLensProviders: { [langauge: string]: vscode.CodeLensProvider[] } =
-  {
-    // python: [new SuggestionsCodeLensProvider(), new PytestCodeLensProvider()],
-    "*": [new SuggestionsCodeLensProvider()],
-  };
+let suggestionsCodeLensDisposable: vscode.Disposable | undefined = undefined;
 
 export function registerAllCodeLensProviders(context: vscode.ExtensionContext) {
-  for (const language in allCodeLensProviders) {
-    for (const codeLensProvider of allCodeLensProviders[language]) {
-      context.subscriptions.push(
-        vscode.languages.registerCodeLensProvider(language, codeLensProvider)
-      );
-    }
+  if (suggestionsCodeLensDisposable) {
+    suggestionsCodeLensDisposable.dispose();
   }
+  suggestionsCodeLensDisposable = vscode.languages.registerCodeLensProvider(
+    "*",
+    new SuggestionsCodeLensProvider()
+  );
+  context.subscriptions.push(suggestionsCodeLensDisposable);
 }
diff --git a/extension/src/suggestions.ts b/extension/src/suggestions.ts
index e269f38a..baa49711 100644
--- a/extension/src/suggestions.ts
+++ b/extension/src/suggestions.ts
@@ -4,6 +4,8 @@ import { openEditorAndRevealRange } from "./util/vscode";
 import { translate, readFileAtRange } from "./util/vscode";
 import * as fs from "fs";
 import * as path from "path";
+import { registerAllCodeLensProviders } from "./lang-server/codeLens";
+import { extensionContext } from "./activation/activate";
 
 export interface SuggestionRanges {
   oldRange: vscode.Range;
@@ -125,6 +127,10 @@ export function rerenderDecorations(editorUri: string) {
     suggestions[idx].newRange,
     vscode.TextEditorRevealType.Default
   );
+
+  if (extensionContext) {
+    registerAllCodeLensProviders(extensionContext);
+  }
 }
 
 export function suggestionDownCommand() {
@@ -337,42 +343,14 @@ export async function showSuggestion(
   range: vscode.Range,
   suggestion: string
 ): Promise<boolean> {
-  // const existingCode = await readFileAtRange(
-  //   new vscode.Range(range.start, range.end),
-  //   editorFilename
-  // );
-
-  // If any of the outside lines are the same, don't repeat them in the suggestion
-  // const slines = suggestion.split("\n");
-  // const elines = existingCode.split("\n");
-  // let linesRemovedBefore = 0;
-  // let linesRemovedAfter = 0;
-  // while (slines.length > 0 && elines.length > 0 && slines[0] === elines[0]) {
-  //   slines.shift();
-  //   elines.shift();
-  //   linesRemovedBefore++;
-  // }
-
-  // while (
-  //   slines.length > 0 &&
-  //   elines.length > 0 &&
-  //   slines[slines.length - 1] === elines[elines.length - 1]
-  // ) {
-  //   slines.pop();
-  //   elines.pop();
-  //   linesRemovedAfter++;
-  // }
-
-  // suggestion = slines.join("\n");
-  // if (suggestion === "") return Promise.resolve(false); // Don't even make a suggestion if they are exactly the same
-
-  // range = new vscode.Range(
-  //   new vscode.Position(range.start.line + linesRemovedBefore, 0),
-  //   new vscode.Position(
-  //     range.end.line - linesRemovedAfter,
-  //     elines.at(-1)?.length || 0
-  //   )
-  // );
+  // Check for empty suggestions:
+  if (
+    suggestion === "" &&
+    range.start.line === range.end.line &&
+    range.start.character === range.end.character
+  ) {
+    return Promise.resolve(false);
+  }
 
   const editor = await openEditorAndRevealRange(editorFilename, range);
   if (!editor) return Promise.resolve(false);
diff --git a/extension/src/util/messenger.ts b/extension/src/util/messenger.ts
index e4133230..b1df161b 100644
--- a/extension/src/util/messenger.ts
+++ b/extension/src/util/messenger.ts
@@ -1,5 +1,6 @@
 console.log("Websocket import");
 const WebSocket = require("ws");
+import fetch from "node-fetch";
 
 export abstract class Messenger {
   abstract send(messageType: string, data: object): void;
@@ -50,18 +51,49 @@ export class WebsocketMessenger extends Messenger {
     return newWebsocket;
   }
 
+  async checkServerRunning(serverUrl: string): Promise<boolean> {
+    // Check if already running by calling /health
+    try {
+      const response = await fetch(serverUrl + "/health");
+      if (response.status === 200) {
+        console.log("Continue python server already running");
+        return true;
+      } else {
+        return false;
+      }
+    } catch (e) {
+      return false;
+    }
+  }
+
   constructor(serverUrl: string) {
     super();
     this.serverUrl = serverUrl;
     this.websocket = this._newWebsocket();
 
-    const interval = setInterval(() => {
-      if (this.websocket.readyState === this.websocket.OPEN) {
-        clearInterval(interval);
-      } else if (this.websocket.readyState !== this.websocket.CONNECTING) {
-        this.websocket = this._newWebsocket();
-      }
-    }, 1000);
+    // Wait until the server is running
+    // const interval = setInterval(async () => {
+    //   if (
+    //     await this.checkServerRunning(
+    //       serverUrl.replace("/ide/ws", "").replace("ws://", "http://")
+    //     )
+    //   ) {
+    //     this.websocket = this._newWebsocket();
+    //     clearInterval(interval);
+    //   } else {
+    //     console.log(
+    //       "Waiting for python server to start-----------------------"
+    //     );
+    //   }
+    // }, 1000);
+
+    // const interval = setInterval(() => {
+    //   if (this.websocket.readyState === this.websocket.OPEN) {
+    //     clearInterval(interval);
+    //   } else if (this.websocket.readyState !== this.websocket.CONNECTING) {
+    //     this.websocket = this._newWebsocket();
+    //   }
+    // }, 1000);
   }
 
   send(messageType: string, data: object) {
@@ -99,10 +131,16 @@ export class WebsocketMessenger extends Messenger {
     });
   }
 
-  onMessage(callback: (messageType: string, data: any) => void): void {
+  onMessage(
+    callback: (
+      messageType: string,
+      data: any,
+      messenger: WebsocketMessenger
+    ) => void
+  ): void {
     this.websocket.addEventListener("message", (event) => {
       const msg = JSON.parse(event.data);
-      callback(msg.messageType, msg.data);
+      callback(msg.messageType, msg.data, this);
     });
   }