3 files changed, 24 insertions, 8 deletions
diff --git a/extension/src/activation/environmentSetup.ts b/extension/src/activation/environmentSetup.ts
index 823670fd..b8c23733 100644
--- a/extension/src/activation/environmentSetup.ts
+++ b/extension/src/activation/environmentSetup.ts
@@ -150,11 +150,10 @@ async function setupPythonEnv() {
     // Repeat until it is successfully created (sometimes it fails to generate the bin, need to try again)
     while (true) {
       const [, stderr] = await runCommand(createEnvCommand);
-      if (stderr) {
-        throw new Error(stderr);
-      }
       if (checkEnvExists()) {
         break;
+      } else if (stderr) {
+        throw new Error(stderr);
       } else {
         // Remove the env and try again
         const removeCommand = `rm -rf "${path.join(
@@ -180,7 +179,6 @@ async function setupPythonEnv() {
         activateCmd,
         pipUpgradeCmd,
         `${pipCmd} install -r requirements.txt`,
-        touchCmd,
       ].join(" ; ");
       const [, stderr] = await runCommand(installRequirementsCommand);
       if (stderr) {
@@ -273,13 +271,14 @@ export async function startContinuePythonServer() {
           console.log(`stdout: ${data}`);
         });
         child.stderr.on("data", (data: any) => {
-          console.log(`stderr: ${data}`);
           if (
             data.includes("Uvicorn running on") || // Successfully started the server
             data.includes("address already in use") // The server is already running (probably a simultaneously opened VS Code window)
           ) {
             console.log("Successfully started Continue python server");
             resolve(null);
+          } else {
+            console.log(`stderr: ${data}`);
           }
         });
         child.on("error", (error: any) => {
diff --git a/extension/src/continueIdeClient.ts b/extension/src/continueIdeClient.ts
index 8ab3e075..d983d93a 100644
--- a/extension/src/continueIdeClient.ts
+++ b/extension/src/continueIdeClient.ts
@@ -226,9 +226,18 @@ class IdeProtocolClient {
   }
 
   async getSessionId(): Promise<string> {
-    if (this.messenger === null) {
-      console.log("MESSENGER IS NULL");
-    }
+    await new Promise((resolve, reject) => {
+      // Repeatedly try to connect to the server
+      const interval = setInterval(() => {
+        if (
+          this.messenger &&
+          this.messenger.websocket.readyState === 1 // 1 => OPEN
+        ) {
+          clearInterval(interval);
+          resolve(null);
+        }
+      }, 1000);
+    });
     const resp = await this.messenger?.sendAndReceive("openGUI", {});
     const sessionId = resp.sessionId;
     console.log("New Continue session with ID: ", sessionId);
diff --git a/extension/src/util/messenger.ts b/extension/src/util/messenger.ts
index 6f8bb29d..e4133230 100644
--- a/extension/src/util/messenger.ts
+++ b/extension/src/util/messenger.ts
@@ -54,6 +54,14 @@ export class WebsocketMessenger extends Messenger {
     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);
   }
 
   send(messageType: string, data: object) {