checkpoint! protocol reform and it works now

8 files changed, 249 insertions, 180 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;
diff --git a/extension/src/commands.ts b/extension/src/commands.ts
index 18f08e31..aeeb4b4f 100644
--- a/extension/src/commands.ts
+++ b/extension/src/commands.ts
@@ -62,11 +62,11 @@ const commandsMap: { [command: string]: (...args: any) => any } = {
   "continue.acceptSuggestion": acceptSuggestionCommand,
   "continue.rejectSuggestion": rejectSuggestionCommand,
   "continue.openDebugPanel": () => {
-    ideProtocolClient?.openNotebook();
+    ideProtocolClient.openNotebook();
   },
   "continue.focusContinueInput": async () => {
     if (!debugPanelWebview) {
-      await ideProtocolClient?.openNotebook();
+      await ideProtocolClient.openNotebook();
     }
     debugPanelWebview?.postMessage({
       type: "focusContinueInput",
diff --git a/extension/src/continueIdeClient.ts b/extension/src/continueIdeClient.ts
index 35eb668d..477d1420 100644
--- a/extension/src/continueIdeClient.ts
+++ b/extension/src/continueIdeClient.ts
@@ -10,30 +10,28 @@ import {
 } from "./suggestions";
 import { debugPanelWebview, setupDebugPanel } from "./debugPanel";
 import { FileEditWithFullContents } from "../schema/FileEditWithFullContents";
-const util = require("util");
-const exec = util.promisify(require("child_process").exec);
-const WebSocket = require("ws");
 import fs = require("fs");
+import { WebsocketMessenger } from "./util/messenger";
 
 class IdeProtocolClient {
-  private _ws: WebSocket | null = null;
-  private _panels: Map<string, vscode.WebviewPanel> = new Map();
-  private readonly _serverUrl: string;
-  private readonly _context: vscode.ExtensionContext;
+  private messenger: WebsocketMessenger | null = null;
+  private panels: Map<string, vscode.WebviewPanel> = new Map();
+  private readonly context: vscode.ExtensionContext;
 
   private _makingEdit = 0;
 
   constructor(serverUrl: string, context: vscode.ExtensionContext) {
-    this._context = context;
-    this._serverUrl = serverUrl;
-    let ws = new WebSocket(serverUrl);
-    this._ws = ws;
-    ws.onclose = () => {
-      this._ws = null;
-    };
-    ws.on("message", (data: any) => {
-      this.handleMessage(JSON.parse(data));
+    this.context = context;
+
+    let messenger = new WebsocketMessenger(serverUrl);
+    this.messenger = messenger;
+    messenger.onClose(() => {
+      this.messenger = null;
+    });
+    messenger.onMessage((messageType, data) => {
+      this.handleMessage(messageType, data);
     });
+
     // Setup listeners for any file changes in open editors
     vscode.workspace.onDidChangeTextDocument((event) => {
       if (this._makingEdit === 0) {
@@ -58,125 +56,52 @@ class IdeProtocolClient {
             };
           }
         );
-        this.send("fileEdits", { fileEdits });
+        this.messenger?.send("fileEdits", { fileEdits });
       } else {
         this._makingEdit--;
       }
     });
   }
 
-  async isConnected() {
-    if (this._ws === null || this._ws.readyState !== WebSocket.OPEN) {
-      let ws = new WebSocket(this._serverUrl);
-      ws.onclose = () => {
-        this._ws = null;
-      };
-      ws.on("message", (data: any) => {
-        this.handleMessage(JSON.parse(data));
-      });
-      this._ws = ws;
-
-      return new Promise((resolve, reject) => {
-        ws.addEventListener("open", () => {
-          resolve(null);
-        });
-      });
-    }
-  }
-
-  async startCore() {
-    var { stdout, stderr } = await exec(
-      "cd /Users/natesesti/Desktop/continue/continue && poetry shell"
-    );
-    if (stderr) {
-      throw new Error(stderr);
-    }
-    var { stdout, stderr } = await exec(
-      "cd .. && uvicorn continue.src.server.main:app --reload --reload-dir continue"
-    );
-    if (stderr) {
-      throw new Error(stderr);
-    }
-    var { stdout, stderr } = await exec("python3 -m continue.src.libs.ide");
-    if (stderr) {
-      throw new Error(stderr);
-    }
-  }
-
-  async send(messageType: string, data: object) {
-    await this.isConnected();
-    let msg = JSON.stringify({ messageType, ...data });
-    this._ws!.send(msg);
-    console.log("Sent message", msg);
-  }
-
-  async receiveMessage(messageType: string): Promise<any> {
-    await this.isConnected();
-    console.log("Connected to websocket");
-    return await new Promise((resolve, reject) => {
-      if (!this._ws) {
-        reject("Not connected to websocket");
-      }
-      this._ws!.onmessage = (event: any) => {
-        let message = JSON.parse(event.data);
-        console.log("RECEIVED MESSAGE", message);
-        if (message.messageType === messageType) {
-          resolve(message);
-        }
-      };
-    });
-  }
-
-  async sendAndReceive(message: any, messageType: string): Promise<any> {
-    try {
-      await this.send(messageType, message);
-      let msg = await this.receiveMessage(messageType);
-      console.log("Received message", msg);
-      return msg;
-    } catch (e) {
-      console.log("Error sending message", e);
-    }
-  }
-
-  async handleMessage(message: any) {
-    switch (message.messageType) {
+  async handleMessage(messageType: string, data: any) {
+    switch (messageType) {
       case "highlightedCode":
-        this.send("highlightedCode", {
+        this.messenger?.send("highlightedCode", {
           highlightedCode: this.getHighlightedCode(),
         });
         break;
       case "workspaceDirectory":
-        this.send("workspaceDirectory", {
+        this.messenger?.send("workspaceDirectory", {
           workspaceDirectory: this.getWorkspaceDirectory(),
         });
       case "openFiles":
-        this.send("openFiles", {
+        this.messenger?.send("openFiles", {
           openFiles: this.getOpenFiles(),
         });
         break;
       case "readFile":
-        this.send("readFile", {
-          contents: this.readFile(message.filepath),
+        this.messenger?.send("readFile", {
+          contents: this.readFile(data.filepath),
         });
         break;
       case "editFile":
-        let fileEdit = await this.editFile(message.edit);
-        this.send("editFile", {
+        const fileEdit = await this.editFile(data.edit);
+        this.messenger?.send("editFile", {
           fileEdit,
         });
         break;
       case "saveFile":
-        this.saveFile(message.filepath);
+        this.saveFile(data.filepath);
         break;
       case "setFileOpen":
-        this.openFile(message.filepath);
+        this.openFile(data.filepath);
         // TODO: Close file
         break;
       case "openNotebook":
       case "connected":
         break;
       default:
-        throw Error("Unknown message type:" + message.messageType);
+        throw Error("Unknown message type:" + messageType);
     }
   }
   getWorkspaceDirectory() {
@@ -209,17 +134,20 @@ class IdeProtocolClient {
   // Initiate Request
 
   closeNotebook(sessionId: string) {
-    this._panels.get(sessionId)?.dispose();
-    this._panels.delete(sessionId);
+    this.panels.get(sessionId)?.dispose();
+    this.panels.delete(sessionId);
   }
 
   async openNotebook() {
     console.log("OPENING NOTEBOOK");
-    let resp = await this.sendAndReceive({}, "openNotebook");
-    let sessionId = resp.sessionId;
+    if (this.messenger === null) {
+      console.log("MESSENGER IS NULL");
+    }
+    const resp = await this.messenger?.sendAndReceive("openNotebook", {});
+    const sessionId = resp.sessionId;
     console.log("SESSION ID", sessionId);
 
-    let column = getRightViewColumn();
+    const column = getRightViewColumn();
     const panel = vscode.window.createWebviewPanel(
       "continue.debugPanelView",
       "Continue",
@@ -231,9 +159,9 @@ class IdeProtocolClient {
     );
 
     // And set its HTML content
-    panel.webview.html = setupDebugPanel(panel, this._context, sessionId);
+    panel.webview.html = setupDebugPanel(panel, this.context, sessionId);
 
-    this._panels.set(sessionId, panel);
+    this.panels.set(sessionId, panel);
   }
 
   acceptRejectSuggestion(accept: boolean, key: SuggestionRanges) {
diff --git a/extension/src/debugPanel.ts b/extension/src/debugPanel.ts
index 4192595c..a295085f 100644
--- a/extension/src/debugPanel.ts
+++ b/extension/src/debugPanel.ts
@@ -16,7 +16,6 @@ import {
 import { sendTelemetryEvent, TelemetryEvent } from "./telemetry";
 import { RangeInFile, SerializedDebugContext } from "./client";
 import { addFileSystemToDebugContext } from "./util/util";
-const WebSocket = require("ws");
 
 class StreamManager {
   private _fullText: string = "";
@@ -108,15 +107,15 @@ class WebsocketConnection {
     this._onOpen = onOpen;
     this._onClose = onClose;
 
-    this._ws.onmessage = (event) => {
+    this._ws.addEventListener("message", (event) => {
       this._onMessage(event.data);
-    };
-    this._ws.onclose = () => {
+    });
+    this._ws.addEventListener("close", () => {
       this._onClose();
-    };
-    this._ws.onopen = () => {
+    });
+    this._ws.addEventListener("open", () => {
       this._onOpen();
-    };
+    });
   }
 
   public send(message: string) {
@@ -230,6 +229,19 @@ export function setupDebugPanel(
           apiUrl: getContinueServerUrl(),
           sessionId,
         });
+
+        // // Listen for changes to server URL in settings
+        // vscode.workspace.onDidChangeConfiguration((event) => {
+        //   if (event.affectsConfiguration("continue.serverUrl")) {
+        //     debugPanelWebview?.postMessage({
+        //       type: "onLoad",
+        //       vscMachineId: vscode.env.machineId,
+        //       apiUrl: getContinueServerUrl(),
+        //       sessionId,
+        //     });
+        //   }
+        // });
+
         break;
       }
 
diff --git a/extension/src/extension.ts b/extension/src/extension.ts
index e0b94278..88af0d19 100644
--- a/extension/src/extension.ts
+++ b/extension/src/extension.ts
@@ -4,7 +4,6 @@
 
 import * as vscode from "vscode";
 import {
-  setupExtensionEnvironment,
   isPythonEnvSetup,
   startContinuePythonServer,
 } from "./activation/environmentSetup";
@@ -26,11 +25,7 @@ export function activate(context: vscode.ExtensionContext) {
       cancellable: false,
     },
     async () => {
-      if (isPythonEnvSetup()) {
-        await startContinuePythonServer();
-      } else {
-        await setupExtensionEnvironment();
-      }
+      await startContinuePythonServer();
       dynamicImportAndActivate(context, true);
     }
   );
diff --git a/extension/src/test/runTest.ts b/extension/src/test/runTest.ts
index 27b3ceb2..e810ed5b 100644
--- a/extension/src/test/runTest.ts
+++ b/extension/src/test/runTest.ts
@@ -1,23 +1,23 @@
-import * as path from 'path';
+import * as path from "path";
 
-import { runTests } from '@vscode/test-electron';
+import { runTests } from "@vscode/test-electron";
 
 async function main() {
-	try {
-		// The folder containing the Extension Manifest package.json
-		// Passed to `--extensionDevelopmentPath`
-		const extensionDevelopmentPath = path.resolve(__dirname, '../../');
+  try {
+    // The folder containing the Extension Manifest package.json
+    // Passed to `--extensionDevelopmentPath`
+    const extensionDevelopmentPath = path.resolve(__dirname, "../../");
 
-		// The path to test runner
-		// Passed to --extensionTestsPath
-		const extensionTestsPath = path.resolve(__dirname, './suite/index');
+    // The path to test runner
+    // Passed to --extensionTestsPath
+    const extensionTestsPath = path.resolve(__dirname, "./suite/index");
 
-		// Download VS Code, unzip it and run the integration test
-		await runTests({ extensionDevelopmentPath, extensionTestsPath });
-	} catch (err) {
-		console.error('Failed to run tests');
-		process.exit(1);
-	}
+    // Download VS Code, unzip it and run the integration test
+    await runTests({ extensionDevelopmentPath, extensionTestsPath });
+  } catch (err) {
+    console.error("Failed to run tests");
+    process.exit(1);
+  }
 }
 
 main();
diff --git a/extension/src/util/messenger.ts b/extension/src/util/messenger.ts
new file mode 100644
index 00000000..6f8bb29d
--- /dev/null
+++ b/extension/src/util/messenger.ts
@@ -0,0 +1,108 @@
+console.log("Websocket import");
+const WebSocket = require("ws");
+
+export abstract class Messenger {
+  abstract send(messageType: string, data: object): void;
+
+  abstract onMessageType(
+    messageType: string,
+    callback: (data: object) => void
+  ): void;
+
+  abstract onMessage(callback: (messageType: string, data: any) => void): void;
+
+  abstract onOpen(callback: () => void): void;
+
+  abstract onClose(callback: () => void): void;
+
+  abstract sendAndReceive(messageType: string, data: any): Promise<any>;
+}
+
+export class WebsocketMessenger extends Messenger {
+  websocket: WebSocket;
+  private onMessageListeners: {
+    [messageType: string]: ((data: object) => void)[];
+  } = {};
+  private onOpenListeners: (() => void)[] = [];
+  private onCloseListeners: (() => void)[] = [];
+  private serverUrl: string;
+
+  _newWebsocket(): WebSocket {
+    // // Dynamic import, because WebSocket is builtin with browser, but not with node. And can't use require in browser.
+    // if (typeof process === "object") {
+    //   console.log("Using node");
+    //   // process is only available in Node
+    //   var WebSocket = require("ws");
+    // }
+
+    const newWebsocket = new WebSocket(this.serverUrl);
+    for (const listener of this.onOpenListeners) {
+      this.onOpen(listener);
+    }
+    for (const listener of this.onCloseListeners) {
+      this.onClose(listener);
+    }
+    for (const messageType in this.onMessageListeners) {
+      for (const listener of this.onMessageListeners[messageType]) {
+        this.onMessageType(messageType, listener);
+      }
+    }
+    return newWebsocket;
+  }
+
+  constructor(serverUrl: string) {
+    super();
+    this.serverUrl = serverUrl;
+    this.websocket = this._newWebsocket();
+  }
+
+  send(messageType: string, data: object) {
+    const payload = JSON.stringify({ messageType, data });
+    if (this.websocket.readyState === this.websocket.OPEN) {
+      this.websocket.send(payload);
+    } else {
+      if (this.websocket.readyState !== this.websocket.CONNECTING) {
+        this.websocket = this._newWebsocket();
+      }
+      this.websocket.addEventListener("open", () => {
+        this.websocket.send(payload);
+      });
+    }
+  }
+
+  sendAndReceive(messageType: string, data: any): Promise<any> {
+    return new Promise((resolve, reject) => {
+      const eventListener = (data: any) => {
+        // THIS ISN"T GETTING CALLED
+        resolve(data);
+        this.websocket.removeEventListener("message", eventListener);
+      };
+      this.onMessageType(messageType, eventListener);
+      this.send(messageType, data);
+    });
+  }
+
+  onMessageType(messageType: string, callback: (data: any) => void): void {
+    this.websocket.addEventListener("message", (event: any) => {
+      const msg = JSON.parse(event.data);
+      if (msg.messageType === messageType) {
+        callback(msg.data);
+      }
+    });
+  }
+
+  onMessage(callback: (messageType: string, data: any) => void): void {
+    this.websocket.addEventListener("message", (event) => {
+      const msg = JSON.parse(event.data);
+      callback(msg.messageType, msg.data);
+    });
+  }
+
+  onOpen(callback: () => void): void {
+    this.websocket.addEventListener("open", callback);
+  }
+
+  onClose(callback: () => void): void {
+    this.websocket.addEventListener("close", callback);
+  }
+}