2 files changed, 67 insertions, 41 deletions

diff --git a/extension/src/activation/activate.ts b/extension/src/activation/activate.ts
index a0aa560b..712ffe13 100644
--- a/extension/src/activation/activate.ts
+++ b/extension/src/activation/activate.ts
@@ -10,7 +10,7 @@ import { getContinueServerUrl } from "../bridge";
 
 export let extensionContext: vscode.ExtensionContext | undefined = undefined;
 
-export let ideProtocolClient: IdeProtocolClient | undefined = undefined;
+export let ideProtocolClient: IdeProtocolClient;
 
 export function activateExtension(
   context: vscode.ExtensionContext,
@@ -24,7 +24,7 @@ export function activateExtension(
   let serverUrl = getContinueServerUrl();
 
   ideProtocolClient = new IdeProtocolClient(
-    serverUrl.replace("http", "ws") + "/ide/ws",
+    `${serverUrl.replace("http", "ws")}/ide/ws`,
     context
   );
 
diff --git a/extension/src/activation/environmentSetup.ts b/extension/src/activation/environmentSetup.ts
index 93a471ff..ad6ac71b 100644
--- a/extension/src/activation/environmentSetup.ts
+++ b/extension/src/activation/environmentSetup.ts
@@ -28,18 +28,7 @@ async function runCommand(cmd: string): Promise<[string, string | undefined]> {
   return [stdout, stderr];
 }
 
-async function getPythonCmdAssumingInstalled() {
-  const [, stderr] = await runCommand("python3 --version");
-  if (stderr) {
-    return "python";
-  }
-  return "python3";
-}
-
-async function setupPythonEnv() {
-  console.log("Setting up python env for Continue extension...");
-  // First check that python3 is installed
-
+async function getPythonPipCommands() {
   var [stdout, stderr] = await runCommand("python3 --version");
   let pythonCmd = "python3";
   if (stderr) {
@@ -58,28 +47,77 @@ async function setupPythonEnv() {
     }
   }
   let pipCmd = pythonCmd.endsWith("3") ? "pip3" : "pip";
+  return [pythonCmd, pipCmd];
+}
 
+function getActivateUpgradeCommands(pythonCmd: string, pipCmd: string) {
   let activateCmd = ". env/bin/activate";
   let pipUpgradeCmd = `${pipCmd} install --upgrade pip`;
   if (process.platform == "win32") {
     activateCmd = ".\\env\\Scripts\\activate";
     pipUpgradeCmd = `${pythonCmd} -m pip install --upgrade pip`;
   }
+  return [activateCmd, pipUpgradeCmd];
+}
 
-  let command = `cd ${path.join(
+function checkEnvExists() {
+  const envBinActivatePath = path.join(
     getExtensionUri().fsPath,
-    "scripts"
-  )} && ${pythonCmd} -m venv env && ${activateCmd} && ${pipUpgradeCmd} && ${pipCmd} install -r requirements.txt`;
-  var [stdout, stderr] = await runCommand(command);
-  if (stderr) {
-    throw new Error(stderr);
+    "scripts",
+    "env",
+    "bin",
+    "activate"
+  );
+  return fs.existsSync(envBinActivatePath);
+}
+
+async function setupPythonEnv() {
+  console.log("Setting up python env for Continue extension...");
+
+  // 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);
+    }
   }
   console.log(
     "Successfully set up python env at ",
     getExtensionUri().fsPath + "/scripts/env"
   );
 
-  await startContinuePythonServer();
+  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);
+  }
 }
 
 function readEnvFile(path: string) {
@@ -116,29 +154,19 @@ function writeEnvFile(path: string, key: string, value: string) {
 }
 
 export async function startContinuePythonServer() {
+  await setupPythonEnv();
+
   // Check vscode settings
   let serverUrl = getContinueServerUrl();
   if (serverUrl !== "http://localhost:8000") {
     return;
   }
 
-  let envFile = path.join(getExtensionUri().fsPath, "scripts", ".env");
-  let openai_api_key: string | undefined =
-    readEnvFile(envFile)["OPENAI_API_KEY"];
-  while (typeof openai_api_key === "undefined" || openai_api_key === "") {
-    openai_api_key = await vscode.window.showInputBox({
-      prompt: "Enter your OpenAI API key",
-      placeHolder: "Enter your OpenAI API key",
-    });
-    // Write to .env file
-  }
-  writeEnvFile(envFile, "OPENAI_API_KEY", openai_api_key);
-
   console.log("Starting Continue python server...");
 
   // Check if already running by calling /health
   try {
-    let response = await fetch(serverUrl + "/health");
+    const response = await fetch(serverUrl + "/health");
     if (response.status === 200) {
       console.log("Continue python server already running");
       return;
@@ -152,15 +180,18 @@ export async function startContinuePythonServer() {
     pythonCmd = "python";
   }
 
+  // let command = `cd ${path.join(
+  //   getExtensionUri().fsPath,
+  //   "scripts"
+  // )} && ${activateCmd} && cd env/lib/python3.11/site-packages && ${pythonCmd} -m continuedev.server.main`;
   let command = `cd ${path.join(
     getExtensionUri().fsPath,
     "scripts"
   )} && ${activateCmd} && cd .. && ${pythonCmd} -m scripts.run_continue_server`;
   try {
     // exec(command);
-    let child = spawn(command, {
+    const child = spawn(command, {
       shell: true,
-      detached: true,
     });
     child.stdout.on("data", (data: any) => {
       console.log(`stdout: ${data}`);
@@ -194,11 +225,6 @@ export function isPythonEnvSetup(): boolean {
   return fs.existsSync(path.join(pathToEnvCfg));
 }
 
-export async function setupExtensionEnvironment() {
-  console.log("Setting up environment for Continue extension...");
-  await Promise.all([setupPythonEnv()]);
-}
-
 export async function downloadPython3() {
   // Download python3 and return the command to run it (python or python3)
   let os = process.platform;