Merge branch 'main' into styled-code

1 files changed, 106 insertions, 69 deletions

diff --git a/extension/src/activation/environmentSetup.ts b/extension/src/activation/environmentSetup.ts
index 593b727e..25b6f643 100644
--- a/extension/src/activation/environmentSetup.ts
+++ b/extension/src/activation/environmentSetup.ts
@@ -7,6 +7,25 @@ import * as fs from "fs";
 import rebuild from "@electron/rebuild";
 import { getContinueServerUrl } from "../bridge";
 import fetch from "node-fetch";
+import * as vscode from "vscode";
+
+const MAX_RETRIES = 5;
+async function retryThenFail(
+  fn: () => Promise<any>,
+  retries: number = MAX_RETRIES
+): Promise<any> {
+  try {
+    return await fn();
+  } catch (e) {
+    if (retries > 0) {
+      return await retryThenFail(fn, retries - 1);
+    }
+    vscode.window.showErrorMessage(
+      "Failed to set up Continue extension. Please email nate@continue.dev and we'll get this fixed ASAP!"
+    );
+    throw e;
+  }
+}
 
 async function runCommand(cmd: string): Promise<[string, string | undefined]> {
   console.log("Running command: ", cmd);
@@ -77,55 +96,71 @@ function checkEnvExists() {
   );
 }
 
+function checkRequirementsInstalled() {
+  return fs.existsSync(
+    path.join(getExtensionUri().fsPath, "scripts", ".continue_env_installed")
+  );
+}
+
 async function setupPythonEnv() {
   console.log("Setting up python env for Continue extension...");
 
-  if (checkEnvExists()) return;
-
-  // Assemble the command to create the env
   const [pythonCmd, pipCmd] = await getPythonPipCommands();
   const [activateCmd, pipUpgradeCmd] = getActivateUpgradeCommands(
     pythonCmd,
     pipCmd
   );
-  const createEnvCommand = [
-    `cd ${path.join(getExtensionUri().fsPath, "scripts")}`,
-    `${pythonCmd} -m venv env`,
-  ].join("; ");
 
-  // 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 {
-      // Remove the env and try again
-      const removeCommand = `rm -rf ${path.join(
-        getExtensionUri().fsPath,
-        "scripts",
-        "env"
-      )}`;
-      await runCommand(removeCommand);
+  if (checkEnvExists()) {
+    console.log("Python env already exists, skipping...");
+  } else {
+    // Assemble the command to create the env
+    const createEnvCommand = [
+      `cd ${path.join(getExtensionUri().fsPath, "scripts")}`,
+      `${pythonCmd} -m venv env`,
+    ].join("; ");
+
+    // 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 {
+        // Remove the env and try again
+        const removeCommand = `rm -rf ${path.join(
+          getExtensionUri().fsPath,
+          "scripts",
+          "env"
+        )}`;
+        await runCommand(removeCommand);
+      }
     }
+    console.log(
+      "Successfully set up python env at ",
+      getExtensionUri().fsPath + "/scripts/env"
+    );
   }
-  console.log(
-    "Successfully set up python env at ",
-    getExtensionUri().fsPath + "/scripts/env"
-  );
 
-  const installRequirementsCommand = [
-    `cd ${path.join(getExtensionUri().fsPath, "scripts")}`,
-    activateCmd,
-    pipUpgradeCmd,
-    `${pipCmd} install -r requirements.txt`,
-  ].join(" && ");
-  const [, stderr] = await runCommand(installRequirementsCommand);
-  if (stderr) {
-    throw new Error(stderr);
-  }
+  await retryThenFail(async () => {
+    if (checkRequirementsInstalled()) {
+      console.log("Python requirements already installed, skipping...");
+    } else {
+      const installRequirementsCommand = [
+        `cd ${path.join(getExtensionUri().fsPath, "scripts")}`,
+        activateCmd,
+        pipUpgradeCmd,
+        `${pipCmd} install -r requirements.txt`,
+        "touch .continue_env_installed",
+      ].join(" ; ");
+      const [, stderr] = await runCommand(installRequirementsCommand);
+      if (stderr) {
+        throw new Error(stderr);
+      }
+    }
+  });
 }
 
 function readEnvFile(path: string) {
@@ -180,15 +215,11 @@ export async function startContinuePythonServer() {
   await setupPythonEnv();
 
   // Check vscode settings
-  let serverUrl = getContinueServerUrl();
+  const serverUrl = getContinueServerUrl();
   if (serverUrl !== "http://localhost:8000") {
     return;
   }
 
-  console.log("Starting Continue python server...");
-
-  if (await checkServerRunning(serverUrl)) return;
-
   let activateCmd = ". env/bin/activate";
   let pythonCmd = "python3";
   if (process.platform == "win32") {
@@ -199,35 +230,41 @@ export async function startContinuePythonServer() {
   let command = `cd ${path.join(
     getExtensionUri().fsPath,
     "scripts"
-  )} && ${activateCmd} && cd .. && ${pythonCmd} -m scripts.run_continue_server`;
+  )} ; ${activateCmd} ; cd .. ; ${pythonCmd} -m scripts.run_continue_server`;
+
+  return await retryThenFail(async () => {
+    console.log("Starting Continue python server...");
+
+    if (await checkServerRunning(serverUrl)) return;
 
-  return new Promise(async (resolve, reject) => {
-    try {
-      const child = spawn(command, {
-        shell: true,
-      });
-      child.stdout.on("data", (data: any) => {
-        console.log(`stdout: ${data}`);
-      });
-      child.stderr.on("data", (data: any) => {
-        console.log(`stderr: ${data}`);
-        if (data.includes("Uvicorn running on")) {
-          console.log("Successfully started Continue python server");
+    return new Promise(async (resolve, reject) => {
+      try {
+        const child = spawn(command, {
+          shell: true,
+        });
+        child.stdout.on("data", (data: any) => {
+          console.log(`stdout: ${data}`);
+        });
+        child.stderr.on("data", (data: any) => {
+          console.log(`stderr: ${data}`);
+          if (data.includes("Uvicorn running on")) {
+            console.log("Successfully started Continue python server");
+            resolve(null);
+          }
+        });
+        child.on("error", (error: any) => {
+          console.log(`error: ${error.message}`);
+        });
+      } catch (e) {
+        console.log("Failed to start Continue python server", e);
+        // If failed, check if it's because the server is already running (might have happened just after we checked above)
+        if (await checkServerRunning(serverUrl)) {
           resolve(null);
+        } else {
+          reject();
         }
-      });
-      child.on("error", (error: any) => {
-        console.log(`error: ${error.message}`);
-      });
-    } catch (e) {
-      console.log("Failed to start Continue python server", e);
-      // If failed, check if it's because the server is already running (might have happened just after we checked above)
-      if (await checkServerRunning(serverUrl)) {
-        resolve(null);
-      } else {
-        reject();
       }
-    }
+    });
   });
 }
 
@@ -255,10 +292,10 @@ export async function downloadPython3() {
     throw new Error("python3 not found");
   } else if (os === "linux") {
     command =
-      "sudo apt update && upgrade && sudo apt install python3 python3-pip";
+      "sudo apt update ; upgrade ; sudo apt install python3 python3-pip";
   } else if (os === "win32") {
     command =
-      "wget -O python_installer.exe https://www.python.org/ftp/python/3.11.3/python-3.11.3-amd64.exe && python_installer.exe /quiet InstallAllUsers=1 PrependPath=1 Include_test=0";
+      "wget -O python_installer.exe https://www.python.org/ftp/python/3.11.3/python-3.11.3-amd64.exe ; python_installer.exe /quiet InstallAllUsers=1 PrependPath=1 Include_test=0";
     pythonCmd = "python";
   }