Merge branch 'main' into newer-simpler-stream-algo

1 files changed, 111 insertions, 74 deletions

diff --git a/extension/src/activation/environmentSetup.ts b/extension/src/activation/environmentSetup.ts
index cdaf3ac0..ec0e228d 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);
@@ -52,14 +71,16 @@ async function getPythonPipCommands() {
   return [pythonCmd, pipCmd];
 }
 
-function getActivateUpgradeCommands(pythonCmd: string, pipCmd: string) {
+function getActivateUpgradeTouchCommands(pythonCmd: string, pipCmd: string) {
   let activateCmd = ". env/bin/activate";
   let pipUpgradeCmd = `${pipCmd} install --upgrade pip`;
+  let touchCmd = "touch .continue_env_installed";
   if (process.platform == "win32") {
     activateCmd = ".\\env\\Scripts\\activate";
     pipUpgradeCmd = `${pythonCmd} -m pip install --upgrade pip`;
+    touchCmd = "ni .continue_env_installed -type file";
   }
-  return [activateCmd, pipUpgradeCmd];
+  return [activateCmd, pipUpgradeCmd, touchCmd];
 }
 
 function checkEnvExists() {
@@ -77,55 +98,69 @@ 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("; ");
+  const [activateCmd, pipUpgradeCmd, touchCmd] =
+    getActivateUpgradeTouchCommands(pythonCmd, pipCmd);
 
-  // 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`,
+        touchCmd,
+      ].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") {
@@ -196,38 +227,44 @@ export async function startContinuePythonServer() {
     pythonCmd = "python";
   }
 
-  let command = `cd ${path.join(
+  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...");
 
-  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");
+    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");
+            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();
       }
-    }
+    });
   });
 }