Merge branch 'main' into ggml-server

13 files changed, 640 insertions, 659 deletions

diff --git a/extension/src/activation/activate.ts b/extension/src/activation/activate.ts
index 18650561..5c6ffa02 100644
--- a/extension/src/activation/activate.ts
+++ b/extension/src/activation/activate.ts
@@ -2,24 +2,41 @@ import * as vscode from "vscode";
 import { registerAllCommands } from "../commands";
 import { registerAllCodeLensProviders } from "../lang-server/codeLens";
 import { sendTelemetryEvent, TelemetryEvent } from "../telemetry";
-// import { openCapturedTerminal } from "../terminal/terminalEmulator";
 import IdeProtocolClient from "../continueIdeClient";
 import { getContinueServerUrl } from "../bridge";
-import { CapturedTerminal } from "../terminal/terminalEmulator";
-import { setupDebugPanel, ContinueGUIWebviewViewProvider } from "../debugPanel";
-import { startContinuePythonServer } from "./environmentSetup";
+import { ContinueGUIWebviewViewProvider } from "../debugPanel";
+import {
+  getExtensionVersion,
+  startContinuePythonServer,
+} from "./environmentSetup";
+import fetch from "node-fetch";
+import registerQuickFixProvider from "../lang-server/codeActions";
 // import { CapturedTerminal } from "../terminal/terminalEmulator";
 
+const PACKAGE_JSON_RAW_GITHUB_URL =
+  "https://raw.githubusercontent.com/continuedev/continue/HEAD/extension/package.json";
+
 export let extensionContext: vscode.ExtensionContext | undefined = undefined;
 
 export let ideProtocolClient: IdeProtocolClient;
 
-export async function activateExtension(
-  context: vscode.ExtensionContext,
-  showTutorial: boolean
-) {
+export async function activateExtension(context: vscode.ExtensionContext) {
   extensionContext = context;
 
+  // Before anything else, check whether this is an out-of-date version of the extension
+  // Do so by grabbing the package.json off of the GitHub respository for now.
+  fetch(PACKAGE_JSON_RAW_GITHUB_URL)
+    .then(async (res) => res.json())
+    .then((packageJson) => {
+      if (packageJson.version !== getExtensionVersion()) {
+        vscode.window.showInformationMessage(
+          `You are using an out-of-date version of the Continue extension. Please update to the latest version.`
+        );
+      }
+    })
+    .catch((e) => console.log("Error checking for extension updates: ", e));
+
+  // Start the Python server
   await new Promise((resolve, reject) => {
     vscode.window.withProgress(
       {
@@ -35,19 +52,20 @@ export async function activateExtension(
     );
   });
 
+  // Register commands and providers
   sendTelemetryEvent(TelemetryEvent.ExtensionActivated);
   registerAllCodeLensProviders(context);
   registerAllCommands(context);
+  registerQuickFixProvider();
 
+  // Initialize IDE Protocol Client
   const serverUrl = getContinueServerUrl();
-
   ideProtocolClient = new IdeProtocolClient(
     `${serverUrl.replace("http", "ws")}/ide/ws`,
     context
   );
 
-  // Setup the left panel
-  (async () => {
+  {
     const sessionIdPromise = await ideProtocolClient.getSessionId();
     const provider = new ContinueGUIWebviewViewProvider(sessionIdPromise);
 
@@ -60,10 +78,5 @@ export async function activateExtension(
         }
       )
     );
-  })();
-  // All opened terminals should be replaced by our own terminal
-  // vscode.window.onDidOpenTerminal((terminal) => {});
-
-  // If any terminals are open to start, replace them
-  // vscode.window.terminals.forEach((terminal) => {}
+  }
 }
diff --git a/extension/src/activation/environmentSetup.ts b/extension/src/activation/environmentSetup.ts
index 714080e3..69a3b75a 100644
--- a/extension/src/activation/environmentSetup.ts
+++ b/extension/src/activation/environmentSetup.ts
@@ -7,23 +7,56 @@ import * as fs from "fs";
 import { getContinueServerUrl } from "../bridge";
 import fetch from "node-fetch";
 import * as vscode from "vscode";
+import * as os from "os";
 import fkill from "fkill";
 import { sendTelemetryEvent, TelemetryEvent } from "../telemetry";
 
+const WINDOWS_REMOTE_SIGNED_SCRIPTS_ERROR =
+  "A Python virtual enviroment cannot be activated because running scripts is disabled for this user. In order to use Continue, please enable signed scripts to run with this command in PowerShell: `Set-ExecutionPolicy -ExecutionPolicy RemoteSigned -Scope CurrentUser`, reload VS Code, and then try again.";
+
 const MAX_RETRIES = 3;
 async function retryThenFail(
   fn: () => Promise<any>,
   retries: number = MAX_RETRIES
 ): Promise<any> {
   try {
+    if (retries < MAX_RETRIES && process.platform === "win32") {
+      let [stdout, stderr] = await runCommand("Get-ExecutionPolicy");
+      if (!stdout.includes("RemoteSigned")) {
+        [stdout, stderr] = await runCommand(
+          "Set-ExecutionPolicy -ExecutionPolicy RemoteSigned -Scope CurrentUser"
+        );
+        console.log("Execution policy stdout: ", stdout);
+        console.log("Execution policy stderr: ", stderr);
+      }
+    }
+
     return await fn();
   } catch (e: any) {
     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!"
-    );
+
+    // Show corresponding error message depending on the platform
+    let msg =
+      "Failed to set up Continue extension. Please email nate@continue.dev and we'll get this fixed ASAP!";
+    try {
+      switch (process.platform) {
+        case "win32":
+          msg = WINDOWS_REMOTE_SIGNED_SCRIPTS_ERROR;
+          break;
+        case "darwin":
+          break;
+        case "linux":
+          const [pythonCmd] = await getPythonPipCommands();
+          msg = await getLinuxAptInstallError(pythonCmd);
+          break;
+      }
+    } finally {
+      console.log("After retries, failed to set up Continue extension", msg);
+      vscode.window.showErrorMessage(msg);
+    }
+
     sendTelemetryEvent(TelemetryEvent.ExtensionSetupError, {
       error: e.message,
     });
@@ -127,8 +160,7 @@ function getActivateUpgradeCommands(pythonCmd: string, pipCmd: string) {
 
 function checkEnvExists() {
   const envBinPath = path.join(
-    getExtensionUri().fsPath,
-    "scripts",
+    serverPath(),
     "env",
     process.platform == "win32" ? "Scripts" : "bin"
   );
@@ -140,10 +172,29 @@ function checkEnvExists() {
   );
 }
 
-function checkRequirementsInstalled() {
+async function checkRequirementsInstalled() {
+  // First, check if the requirements have been installed most recently for a later version of the extension
+  if (fs.existsSync(requirementsVersionPath())) {
+    const requirementsVersion = fs.readFileSync(
+      requirementsVersionPath(),
+      "utf8"
+    );
+    if (requirementsVersion !== getExtensionVersion()) {
+      // Remove the old version of continuedev from site-packages
+      const [pythonCmd, pipCmd] = await getPythonPipCommands();
+      const [activateCmd] = getActivateUpgradeCommands(pythonCmd, pipCmd);
+      const removeOldVersionCommand = [
+        `cd "${serverPath()}"`,
+        activateCmd,
+        `${pipCmd} uninstall -y continuedev`,
+      ].join(" ; ");
+      await runCommand(removeOldVersionCommand);
+      return false;
+    }
+  }
+
   let envLibsPath = path.join(
-    getExtensionUri().fsPath,
-    "scripts",
+    serverPath(),
     "env",
     process.platform == "win32" ? "Lib" : "lib"
   );
@@ -165,27 +216,30 @@ function checkRequirementsInstalled() {
   const continuePath = path.join(envLibsPath, "continuedev");
 
   return fs.existsSync(continuePath);
-
-  // return fs.existsSync(
-  //   path.join(getExtensionUri().fsPath, "scripts", ".continue_env_installed")
-  // );
 }
 
-async function setupPythonEnv() {
-  console.log("Setting up python env for Continue extension...");
-
-  const [pythonCmd, pipCmd] = await getPythonPipCommands();
-  const [activateCmd, pipUpgradeCmd] = getActivateUpgradeCommands(
-    pythonCmd,
-    pipCmd
-  );
+async function getLinuxAptInstallError(pythonCmd: string) {
+  // First, try to run the command to install python3-venv
+  let [stdout, stderr] = await runCommand(`${pythonCmd} --version`);
+  if (stderr) {
+    await vscode.window.showErrorMessage(
+      "Python3 is not installed. Please install from https://www.python.org/downloads, reload VS Code, and try again."
+    );
+    throw new Error(stderr);
+  }
+  const version = stdout.split(" ")[1].split(".")[1];
+  const installVenvCommand = `apt-get install python3.${version}-venv`;
+  await runCommand("apt-get update");
+  return `[Important] Continue needs to create a Python virtual environment, but python3.${version}-venv is not installed. Please run this command in your terminal: \`${installVenvCommand}\`, reload VS Code, and then try again.`;
+}
 
+async function createPythonVenv(pythonCmd: string) {
   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")}"`,
+      `cd "${serverPath()}"`,
       `${pythonCmd} -m venv env`,
     ].join(" ; ");
 
@@ -194,31 +248,37 @@ async function setupPythonEnv() {
       stderr &&
       stderr.includes("running scripts is disabled on this system")
     ) {
-      await vscode.window.showErrorMessage(
-        "A Python virtual enviroment cannot be activated because running scripts is disabled for this user. Please enable signed scripts to run with this command in PowerShell: `Set-ExecutionPolicy -ExecutionPolicy RemoteSigned -Scope CurrentUser`, reload VS Code, and then try again."
-      );
+      console.log("Scripts disabled error when trying to create env");
+      await vscode.window.showErrorMessage(WINDOWS_REMOTE_SIGNED_SCRIPTS_ERROR);
       throw new Error(stderr);
     } else if (
       stderr?.includes("On Debian/Ubuntu systems") ||
       stdout?.includes("On Debian/Ubuntu systems")
     ) {
-      // First, try to run the command to install python3-venv
-      let [stdout, stderr] = await runCommand(`${pythonCmd} --version`);
-      if (stderr) {
-        throw new Error(stderr);
-      }
-      const version = stdout.split(" ")[1].split(".")[1];
-      const installVenvCommand = `apt-get install python3.${version}-venv`;
-      await runCommand("apt-get update");
-      // Ask the user to run the command to install python3-venv (requires sudo, so we can't)
-      // First, get the python version
-      const msg = `[Important] Continue needs to create a Python virtual environment, but python3.${version}-venv is not installed. Please run this command in your terminal: \`${installVenvCommand}\`, reload VS Code, and then try again.`;
+      const msg = await getLinuxAptInstallError(pythonCmd);
       console.log(msg);
       await vscode.window.showErrorMessage(msg);
     } else if (checkEnvExists()) {
-      console.log(
-        "Successfully set up python env at ",
-        getExtensionUri().fsPath + "/scripts/env"
+      console.log("Successfully set up python env at ", `${serverPath()}/env`);
+    } else if (
+      stderr?.includes("Permission denied") &&
+      stderr?.includes("python.exe")
+    ) {
+      // This might mean that another window is currently using the python.exe file to install requirements
+      // So we want to wait and try again
+      let i = 0;
+      await new Promise((resolve, reject) =>
+        setInterval(() => {
+          if (i > 5) {
+            reject("Timed out waiting for other window to create env...");
+          }
+          if (checkEnvExists()) {
+            resolve(null);
+          } else {
+            console.log("Waiting for other window to create env...");
+          }
+          i++;
+        }, 5000)
       );
     } else {
       const msg = [
@@ -230,13 +290,27 @@ async function setupPythonEnv() {
       throw new Error(msg);
     }
   }
+}
+
+async function setupPythonEnv() {
+  console.log("Setting up python env for Continue extension...");
+
+  const [pythonCmd, pipCmd] = await getPythonPipCommands();
+  const [activateCmd, pipUpgradeCmd] = getActivateUpgradeCommands(
+    pythonCmd,
+    pipCmd
+  );
 
   await retryThenFail(async () => {
-    if (checkRequirementsInstalled()) {
+    // First, create the virtual environment
+    await createPythonVenv(pythonCmd);
+
+    // Install the requirements
+    if (await checkRequirementsInstalled()) {
       console.log("Python requirements already installed, skipping...");
     } else {
       const installRequirementsCommand = [
-        `cd "${path.join(getExtensionUri().fsPath, "scripts")}"`,
+        `cd "${serverPath()}"`,
         activateCmd,
         pipUpgradeCmd,
         `${pipCmd} install -r requirements.txt`,
@@ -245,6 +319,8 @@ async function setupPythonEnv() {
       if (stderr) {
         throw new Error(stderr);
       }
+      // Write the version number for which requirements were installed
+      fs.writeFileSync(requirementsVersionPath(), getExtensionVersion());
     }
   });
 }
@@ -297,12 +373,50 @@ async function checkServerRunning(serverUrl: string): Promise<boolean> {
   }
 }
 
+export function getContinueGlobalPath(): string {
+  // This is ~/.continue on mac/linux
+  const continuePath = path.join(os.homedir(), ".continue");
+  if (!fs.existsSync(continuePath)) {
+    fs.mkdirSync(continuePath);
+  }
+  return continuePath;
+}
+
+function setupServerPath() {
+  const sPath = serverPath();
+  const extensionServerPath = path.join(getExtensionUri().fsPath, "server");
+  const files = fs.readdirSync(extensionServerPath);
+  files.forEach((file) => {
+    const filePath = path.join(extensionServerPath, file);
+    fs.copyFileSync(filePath, path.join(sPath, file));
+  });
+}
+
+function serverPath(): string {
+  const sPath = path.join(getContinueGlobalPath(), "server");
+  if (!fs.existsSync(sPath)) {
+    fs.mkdirSync(sPath);
+  }
+  return sPath;
+}
+
+export function devDataPath(): string {
+  const sPath = path.join(getContinueGlobalPath(), "dev_data");
+  if (!fs.existsSync(sPath)) {
+    fs.mkdirSync(sPath);
+  }
+  return sPath;
+}
+
 function serverVersionPath(): string {
-  const extensionPath = getExtensionUri().fsPath;
-  return path.join(extensionPath, "server_version.txt");
+  return path.join(serverPath(), "server_version.txt");
 }
 
-function getExtensionVersion() {
+function requirementsVersionPath(): string {
+  return path.join(serverPath(), "requirements_version.txt");
+}
+
+export function getExtensionVersion() {
   const extension = vscode.extensions.getExtension("continue.continue");
   return extension?.packageJSON.version || "";
 }
@@ -314,39 +428,40 @@ export async function startContinuePythonServer() {
     return;
   }
 
+  setupServerPath();
+
   return await retryThenFail(async () => {
-    if (await checkServerRunning(serverUrl)) {
-      // Kill the server if it is running an old version
-      if (fs.existsSync(serverVersionPath())) {
-        const serverVersion = fs.readFileSync(serverVersionPath(), "utf8");
-        if (serverVersion === getExtensionVersion()) {
-          return;
-        }
+    // Kill the server if it is running an old version
+    if (fs.existsSync(serverVersionPath())) {
+      const serverVersion = fs.readFileSync(serverVersionPath(), "utf8");
+      if (
+        serverVersion === getExtensionVersion() &&
+        (await checkServerRunning(serverUrl))
+      ) {
+        // The current version is already up and running, no need to continue
+        return;
       }
-      console.log("Killing old server...");
-      try {
-        await fkill(":65432");
-      } catch (e) {
-        console.log(
-          "Failed to kill old server, likely because it didn't exist:",
-          e
-        );
+    }
+    console.log("Killing old server...");
+    try {
+      await fkill(":65432");
+    } catch (e: any) {
+      if (!e.message.includes("Process doesn't exist")) {
+        console.log("Failed to kill old server:", e);
       }
     }
 
     // Do this after above check so we don't have to waste time setting up the env
     await setupPythonEnv();
 
+    // Spawn the server process on port 65432
     const [pythonCmd] = await getPythonPipCommands();
     const activateCmd =
       process.platform == "win32"
         ? ".\\env\\Scripts\\activate"
         : ". env/bin/activate";
 
-    const command = `cd "${path.join(
-      getExtensionUri().fsPath,
-      "scripts"
-    )}" && ${activateCmd} && cd .. && ${pythonCmd} -m scripts.run_continue_server`;
+    const command = `cd "${serverPath()}" && ${activateCmd} && cd .. && ${pythonCmd} -m server.run_continue_server`;
 
     console.log("Starting Continue python server...");
 
@@ -359,7 +474,8 @@ export async function startContinuePythonServer() {
           console.log(`stdout: ${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)
+            data.includes("only one usage of each socket address") || // [windows] The server is already running (probably a simultaneously opened VS Code window)
+            data.includes("address already in use") // [mac/linux] The server is already running (probably a simultaneously opened VS Code window)
           ) {
             console.log("Successfully started Continue python server");
             resolve(null);
@@ -392,8 +508,8 @@ export async function startContinuePythonServer() {
 }
 
 export function isPythonEnvSetup(): boolean {
-  let pathToEnvCfg = getExtensionUri().fsPath + "/scripts/env/pyvenv.cfg";
-  return fs.existsSync(path.join(pathToEnvCfg));
+  const pathToEnvCfg = path.join(serverPath(), "env", "pyvenv.cfg");
+  return fs.existsSync(pathToEnvCfg);
 }
 
 export async function downloadPython3() {
diff --git a/extension/src/bridge.ts b/extension/src/bridge.ts
index 92ba4044..7e6398be 100644
--- a/extension/src/bridge.ts
+++ b/extension/src/bridge.ts
@@ -4,14 +4,11 @@ import * as vscode from "vscode";
 import {
   Configuration,
   DebugApi,
-  RangeInFile,
-  SerializedDebugContext,
   UnittestApi,
 } from "./client";
 import { convertSingleToDoubleQuoteJSON } from "./util/util";
 import { getExtensionUri } from "./util/vscode";
 import { extensionContext } from "./activation/activate";
-const axios = require("axios").default;
 const util = require("util");
 const exec = util.promisify(require("child_process").exec);
 
@@ -36,21 +33,16 @@ const configuration = new Configuration({
 export const debugApi = new DebugApi(configuration);
 export const unittestApi = new UnittestApi(configuration);
 
-function get_python_path() {
-  return path.join(getExtensionUri().fsPath, "..");
-}
-
 export function get_api_url() {
-  let extensionUri = getExtensionUri();
-  let configFile = path.join(extensionUri.fsPath, "config/config.json");
-  let config = require(configFile);
+  const extensionUri = getExtensionUri();
+  const configFile = path.join(extensionUri.fsPath, "config/config.json");
+  const config = require(configFile);
 
   if (config.API_URL) {
     return config.API_URL;
   }
   return "http://localhost:65432";
 }
-const API_URL = get_api_url();
 
 export function getContinueServerUrl() {
   // If in debug mode, always use 8001
@@ -66,10 +58,6 @@ export function getContinueServerUrl() {
   );
 }
 
-function build_python_command(cmd: string): string {
-  return `cd ${get_python_path()} && source env/bin/activate && ${cmd}`;
-}
-
 function listToCmdLineArgs(list: string[]): string {
   return list.map((el) => `"$(echo "${el}")"`).join(" ");
 }
@@ -103,198 +91,3 @@ export async function runPythonScript(
     }
   }
 }
-
-function parseStdout(
-  stdout: string,
-  key: string,
-  until_end: boolean = false
-): string {
-  const prompt = `${key}=`;
-  let lines = stdout.split("\n");
-
-  let value: string = "";
-  for (let i = 0; i < lines.length; i++) {
-    if (lines[i].startsWith(prompt)) {
-      if (until_end) {
-        return lines.slice(i).join("\n").substring(prompt.length);
-      } else {
-        return lines[i].substring(prompt.length);
-      }
-    }
-  }
-  return "";
-}
-
-export async function askQuestion(
-  question: string,
-  workspacePath: string
-): Promise<{ answer: string; range: vscode.Range; filename: string }> {
-  const command = build_python_command(
-    `python3 ${path.join(
-      get_python_path(),
-      "ask.py"
-    )} ask ${workspacePath} "${question}"`
-  );
-
-  const { stdout, stderr } = await exec(command);
-  if (stderr) {
-    throw new Error(stderr);
-  }
-  // Use the output
-  const answer = parseStdout(stdout, "Answer");
-  const filename = parseStdout(stdout, "Filename");
-  const startLineno = parseInt(parseStdout(stdout, "Start lineno"));
-  const endLineno = parseInt(parseStdout(stdout, "End lineno"));
-  const range = new vscode.Range(
-    new vscode.Position(startLineno, 0),
-    new vscode.Position(endLineno, 0)
-  );
-  if (answer && filename && startLineno && endLineno) {
-    return { answer, filename, range };
-  } else {
-    throw new Error("Error: No answer found");
-  }
-}
-
-export async function apiRequest(
-  endpoint: string,
-  options: {
-    method?: string;
-    query?: { [key: string]: any };
-    body?: { [key: string]: any };
-  }
-): Promise<any> {
-  let defaults = {
-    method: "GET",
-    query: {},
-    body: {},
-  };
-  options = Object.assign(defaults, options); // Second takes over first
-  if (endpoint.startsWith("/")) endpoint = endpoint.substring(1);
-  console.log("API request: ", options.body);
-
-  let resp;
-  try {
-    resp = await axios({
-      method: options.method,
-      url: `${API_URL}/${endpoint}`,
-      data: options.body,
-      params: options.query,
-      headers: {
-        "x-vsc-machine-id": vscode.env.machineId,
-      },
-    });
-  } catch (err) {
-    console.log("Error: ", err);
-    throw err;
-  }
-
-  return resp.data;
-}
-
-// Write a docstring for the most specific function or class at the current line in the given file
-export async function writeDocstringForFunction(
-  filename: string,
-  position: vscode.Position
-): Promise<{ lineno: number; docstring: string }> {
-  let resp = await apiRequest("docstring/forline", {
-    query: {
-      filecontents: (
-        await vscode.workspace.fs.readFile(vscode.Uri.file(filename))
-      ).toString(),
-      lineno: position.line.toString(),
-    },
-  });
-
-  const lineno = resp.lineno;
-  const docstring = resp.completion;
-  if (lineno && docstring) {
-    return { lineno, docstring };
-  } else {
-    throw new Error("Error: No docstring returned");
-  }
-}
-
-export async function findSuspiciousCode(
-  ctx: SerializedDebugContext
-): Promise<RangeInFile[]> {
-  if (!ctx.traceback) return [];
-  let files = await getFileContents(
-    getFilenamesFromPythonStacktrace(ctx.traceback)
-  );
-  let resp = await debugApi.findSusCodeEndpointDebugFindPost({
-    findBody: {
-      traceback: ctx.traceback,
-      description: ctx.description,
-      filesystem: files,
-    },
-  });
-  let ranges = resp.response;
-  if (
-    ranges.length <= 1 &&
-    ctx.traceback &&
-    ctx.traceback.includes("AssertionError")
-  ) {
-    let parsed_traceback =
-      await debugApi.parseTracebackEndpointDebugParseTracebackGet({
-        traceback: ctx.traceback,
-      });
-    let last_frame = parsed_traceback.frames[0];
-    if (!last_frame) return [];
-    ranges = (
-      await runPythonScript("build_call_graph.py", [
-        last_frame.filepath,
-        last_frame.lineno.toString(),
-        last_frame._function,
-      ])
-    ).value;
-  }
-
-  return ranges;
-}
-
-export async function writeUnitTestForFunction(
-  filename: string,
-  position: vscode.Position
-): Promise<string> {
-  let resp = await apiRequest("unittest/forline", {
-    method: "POST",
-    body: {
-      filecontents: (
-        await vscode.workspace.fs.readFile(vscode.Uri.file(filename))
-      ).toString(),
-      lineno: position.line,
-      userid: vscode.env.machineId,
-    },
-  });
-
-  return resp.completion;
-}
-
-async function getFileContents(
-  files: string[]
-): Promise<{ [key: string]: string }> {
-  let contents = await Promise.all(
-    files.map(async (file: string) => {
-      return (
-        await vscode.workspace.fs.readFile(vscode.Uri.file(file))
-      ).toString();
-    })
-  );
-  let fileContents: { [key: string]: string } = {};
-  for (let i = 0; i < files.length; i++) {
-    fileContents[files[i]] = contents[i];
-  }
-  return fileContents;
-}
-
-function getFilenamesFromPythonStacktrace(traceback: string): string[] {
-  let filenames: string[] = [];
-  for (let line of traceback.split("\n")) {
-    let match = line.match(/File "(.*)", line/);
-    if (match) {
-      filenames.push(match[1]);
-    }
-  }
-  return filenames;
-}
diff --git a/extension/src/commands.ts b/extension/src/commands.ts
index ffb67ab5..2b7f4c0c 100644
--- a/extension/src/commands.ts
+++ b/extension/src/commands.ts
@@ -16,40 +16,16 @@ import {
 import { acceptDiffCommand, rejectDiffCommand } from "./diffs";
 import * as bridge from "./bridge";
 import { debugPanelWebview } from "./debugPanel";
-import { sendTelemetryEvent, TelemetryEvent } from "./telemetry";
 import { ideProtocolClient } from "./activation/activate";
 
-// Copy everything over from extension.ts
-const commandsMap: { [command: string]: (...args: any) => any } = {
-  "continue.askQuestion": (data: any, webviewView: vscode.WebviewView) => {
-    if (!vscode.workspace.workspaceFolders) {
-      return;
-    }
-
-    answerQuestion(
-      data.question,
-      vscode.workspace.workspaceFolders[0].uri.fsPath,
-      webviewView.webview
-    );
-  },
-  "continue.askQuestionFromInput": () => {
-    vscode.window
-      .showInputBox({ placeHolder: "Ask away!" })
-      .then((question) => {
-        if (!question || !vscode.workspace.workspaceFolders) {
-          return;
-        }
+let focusedOnContinueInput = false;
 
-        sendTelemetryEvent(TelemetryEvent.UniversalPromptQuery, {
-          query: question,
-        });
+export const setFocusedOnContinueInput = (value: boolean) => {
+  focusedOnContinueInput = value;
+};
 
-        answerQuestion(
-          question,
-          vscode.workspace.workspaceFolders[0].uri.fsPath
-        );
-      });
-  },
+// Copy everything over from extension.ts
+const commandsMap: { [command: string]: (...args: any) => any } = {
   "continue.suggestionDown": suggestionDownCommand,
   "continue.suggestionUp": suggestionUpCommand,
   "continue.acceptSuggestion": acceptSuggestionCommand,
@@ -58,11 +34,23 @@ const commandsMap: { [command: string]: (...args: any) => any } = {
   "continue.rejectDiff": rejectDiffCommand,
   "continue.acceptAllSuggestions": acceptAllSuggestionsCommand,
   "continue.rejectAllSuggestions": rejectAllSuggestionsCommand,
+  "continue.quickFix": async (message: string, code: string, edit: boolean) => {
+    ideProtocolClient.sendMainUserInput(
+      `${
+        edit ? "/edit " : ""
+      }${code}\n\nHow do I fix this problem in the above code?: ${message}`
+    );
+  },
   "continue.focusContinueInput": async () => {
-    vscode.commands.executeCommand("continue.continueGUIView.focus");
-    debugPanelWebview?.postMessage({
-      type: "focusContinueInput",
-    });
+    if (focusedOnContinueInput) {
+      vscode.commands.executeCommand("workbench.action.focusActiveEditorGroup");
+    } else {
+      vscode.commands.executeCommand("continue.continueGUIView.focus");
+      debugPanelWebview?.postMessage({
+        type: "focusContinueInput",
+      });
+    }
+    focusedOnContinueInput = !focusedOnContinueInput;
   },
   "continue.quickTextEntry": async () => {
     const text = await vscode.window.showInputBox({
@@ -73,27 +61,6 @@ const commandsMap: { [command: string]: (...args: any) => any } = {
     if (text) {
       ideProtocolClient.sendMainUserInput(text);
     }
-    vscode.commands.executeCommand("continue.continueGUIView.focus");
-  },
-};
-
-const textEditorCommandsMap: { [command: string]: (...args: any) => {} } = {
-  "continue.writeDocstring": async (editor: vscode.TextEditor, _) => {
-    sendTelemetryEvent(TelemetryEvent.GenerateDocstring);
-    let gutterSpinnerKey = showGutterSpinner(
-      editor,
-      editor.selection.active.line
-    );
-
-    const { lineno, docstring } = await bridge.writeDocstringForFunction(
-      editor.document.fileName,
-      editor.selection.active
-    );
-    // Can't use the edit given above after an async call
-    editor.edit((edit) => {
-      edit.insert(new vscode.Position(lineno, 0), docstring);
-      decorationManager.deleteDecoration(gutterSpinnerKey);
-    });
   },
 };
 
@@ -103,44 +70,4 @@ export function registerAllCommands(context: vscode.ExtensionContext) {
       vscode.commands.registerCommand(command, callback)
     );
   }
-
-  for (const [command, callback] of Object.entries(textEditorCommandsMap)) {
-    context.subscriptions.push(
-      vscode.commands.registerTextEditorCommand(command, callback)
-    );
-  }
-}
-
-async function answerQuestion(
-  question: string,
-  workspacePath: string,
-  webview: vscode.Webview | undefined = undefined
-) {
-  vscode.window.withProgress(
-    {
-      location: vscode.ProgressLocation.Notification,
-      title: "Anwering question...",
-      cancellable: false,
-    },
-    async (progress, token) => {
-      try {
-        let resp = await bridge.askQuestion(question, workspacePath);
-        // Send the answer back to the webview
-        if (webview) {
-          webview.postMessage({
-            type: "answerQuestion",
-            answer: resp.answer,
-          });
-        }
-        showAnswerInTextEditor(resp.filename, resp.range, resp.answer);
-      } catch (error: any) {
-        if (webview) {
-          webview.postMessage({
-            type: "answerQuestion",
-            answer: error,
-          });
-        }
-      }
-    }
-  );
 }
diff --git a/extension/src/continueIdeClient.ts b/extension/src/continueIdeClient.ts
index 679d94ba..fac0a227 100644
--- a/extension/src/continueIdeClient.ts
+++ b/extension/src/continueIdeClient.ts
@@ -1,10 +1,9 @@
-// import { ShowSuggestionRequest } from "../schema/ShowSuggestionRequest";
 import {
   editorSuggestionsLocked,
   showSuggestion as showSuggestionInEditor,
   SuggestionRanges,
 } from "./suggestions";
-import { openEditorAndRevealRange, getRightViewColumn } from "./util/vscode";
+import { openEditorAndRevealRange } from "./util/vscode";
 import { FileEdit } from "../schema/FileEdit";
 import { RangeInFile } from "../schema/RangeInFile";
 import * as vscode from "vscode";
@@ -15,9 +14,8 @@ import {
 import { FileEditWithFullContents } from "../schema/FileEditWithFullContents";
 import fs = require("fs");
 import { WebsocketMessenger } from "./util/messenger";
-import * as path from "path";
-import * as os from "os";
 import { diffManager } from "./diffs";
+import path = require("path");
 
 class IdeProtocolClient {
   private messenger: WebsocketMessenger | null = null;
@@ -27,17 +25,54 @@ class IdeProtocolClient {
 
   private _highlightDebounce: NodeJS.Timeout | null = null;
 
-  constructor(serverUrl: string, context: vscode.ExtensionContext) {
-    this.context = context;
+  private _lastReloadTime: number = 16;
+  private _reconnectionTimeouts: NodeJS.Timeout[] = [];
+
+  private _sessionId: string | null = null;
+  private _serverUrl: string;
 
-    let messenger = new WebsocketMessenger(serverUrl);
+  private _newWebsocketMessenger() {
+    const requestUrl =
+      this._serverUrl +
+      (this._sessionId ? `?session_id=${this._sessionId}` : "");
+    const messenger = new WebsocketMessenger(requestUrl);
     this.messenger = messenger;
-    messenger.onClose(() => {
+
+    const reconnect = () => {
+      console.log("Trying to reconnect IDE protocol websocket...");
       this.messenger = null;
+
+      // Exponential backoff to reconnect
+      this._reconnectionTimeouts.forEach((to) => clearTimeout(to));
+
+      const timeout = setTimeout(() => {
+        if (this.messenger?.websocket?.readyState === 1) {
+          return;
+        }
+        this._newWebsocketMessenger();
+      }, this._lastReloadTime);
+
+      this._reconnectionTimeouts.push(timeout);
+      this._lastReloadTime = Math.min(2 * this._lastReloadTime, 5000);
+    };
+    messenger.onOpen(() => {
+      this._reconnectionTimeouts.forEach((to) => clearTimeout(to));
+    });
+    messenger.onClose(() => {
+      reconnect();
+    });
+    messenger.onError(() => {
+      reconnect();
     });
     messenger.onMessage((messageType, data, messenger) => {
       this.handleMessage(messageType, data, messenger);
     });
+  }
+
+  constructor(serverUrl: string, context: vscode.ExtensionContext) {
+    this.context = context;
+    this._serverUrl = serverUrl;
+    this._newWebsocketMessenger();
 
     // Setup listeners for any file changes in open editors
     // vscode.workspace.onDidChangeTextDocument((event) => {
@@ -69,8 +104,11 @@ class IdeProtocolClient {
     //   }
     // });
 
-    // Setup listeners for any file changes in open editors
+    // Setup listeners for any selection changes in open editors
     vscode.window.onDidChangeTextEditorSelection((event) => {
+      if (this.editorIsTerminal(event.textEditor)) {
+        return;
+      }
       if (this._highlightDebounce) {
         clearTimeout(this._highlightDebounce);
       }
@@ -131,6 +169,11 @@ class IdeProtocolClient {
           openFiles: this.getOpenFiles(),
         });
         break;
+      case "visibleFiles":
+        messenger.send("visibleFiles", {
+          visibleFiles: this.getVisibleFiles(),
+        });
+        break;
       case "readFile":
         messenger.send("readFile", {
           contents: this.readFile(data.filepath),
@@ -166,7 +209,7 @@ class IdeProtocolClient {
       case "showDiff":
         this.showDiff(data.filepath, data.replacement, data.step_index);
         break;
-      case "openGUI":
+      case "getSessionId":
       case "connected":
         break;
       default:
@@ -279,10 +322,6 @@ class IdeProtocolClient {
   // ------------------------------------ //
   // Initiate Request
 
-  async openGUI(asRightWebviewPanel: boolean = false) {
-    // Open the webview panel
-  }
-
   async getSessionId(): Promise<string> {
     await new Promise((resolve, reject) => {
       // Repeatedly try to connect to the server
@@ -298,10 +337,10 @@ class IdeProtocolClient {
         }
       }, 1000);
     });
-    const resp = await this.messenger?.sendAndReceive("openGUI", {});
-    const sessionId = resp.sessionId;
-    console.log("New Continue session with ID: ", sessionId);
-    return sessionId;
+    const resp = await this.messenger?.sendAndReceive("getSessionId", {});
+    // console.log("New Continue session with ID: ", sessionId);
+    this._sessionId = resp.sessionId;
+    return resp.sessionId;
   }
 
   acceptRejectSuggestion(accept: boolean, key: SuggestionRanges) {
@@ -315,38 +354,55 @@ class IdeProtocolClient {
   // ------------------------------------ //
   // Respond to request
 
+  private editorIsTerminal(editor: vscode.TextEditor) {
+    return (
+      !!path.basename(editor.document.uri.fsPath).match(/\d/) ||
+      (editor.document.languageId === "plaintext" &&
+        editor.document.getText() === "accessible-buffer-accessible-buffer-")
+    );
+  }
+
   getOpenFiles(): string[] {
     return vscode.window.visibleTextEditors
-      .filter((editor) => {
-        return !(
-          editor.document.uri.fsPath.endsWith("/1") ||
-          (editor.document.languageId === "plaintext" &&
-            editor.document.getText() ===
-              "accessible-buffer-accessible-buffer-")
-        );
-      })
+      .filter((editor) => !this.editorIsTerminal(editor))
+      .map((editor) => {
+        return editor.document.uri.fsPath;
+      });
+  }
+
+  getVisibleFiles(): string[] {
+    return vscode.window.visibleTextEditors
+      .filter((editor) => !this.editorIsTerminal(editor))
       .map((editor) => {
         return editor.document.uri.fsPath;
       });
   }
 
   saveFile(filepath: string) {
-    vscode.window.visibleTextEditors.forEach((editor) => {
-      if (editor.document.uri.fsPath === filepath) {
-        editor.document.save();
-      }
-    });
+    vscode.window.visibleTextEditors
+      .filter((editor) => !this.editorIsTerminal(editor))
+      .forEach((editor) => {
+        if (editor.document.uri.fsPath === filepath) {
+          editor.document.save();
+        }
+      });
   }
 
   readFile(filepath: string): string {
     let contents: string | undefined;
-    vscode.window.visibleTextEditors.forEach((editor) => {
-      if (editor.document.uri.fsPath === filepath) {
-        contents = editor.document.getText();
+    vscode.window.visibleTextEditors
+      .filter((editor) => !this.editorIsTerminal(editor))
+      .forEach((editor) => {
+        if (editor.document.uri.fsPath === filepath) {
+          contents = editor.document.getText();
+        }
+      });
+    if (typeof contents === "undefined") {
+      if (fs.existsSync(filepath)) {
+        contents = fs.readFileSync(filepath, "utf-8");
+      } else {
+        contents = "";
       }
-    });
-    if (!contents) {
-      contents = fs.readFileSync(filepath, "utf-8");
     }
     return contents;
   }
@@ -380,25 +436,27 @@ class IdeProtocolClient {
   getHighlightedCode(): RangeInFile[] {
     // TODO
     let rangeInFiles: RangeInFile[] = [];
-    vscode.window.visibleTextEditors.forEach((editor) => {
-      editor.selections.forEach((selection) => {
-        // if (!selection.isEmpty) {
-        rangeInFiles.push({
-          filepath: editor.document.uri.fsPath,
-          range: {
-            start: {
-              line: selection.start.line,
-              character: selection.start.character,
-            },
-            end: {
-              line: selection.end.line,
-              character: selection.end.character,
+    vscode.window.visibleTextEditors
+      .filter((editor) => !this.editorIsTerminal(editor))
+      .forEach((editor) => {
+        editor.selections.forEach((selection) => {
+          // if (!selection.isEmpty) {
+          rangeInFiles.push({
+            filepath: editor.document.uri.fsPath,
+            range: {
+              start: {
+                line: selection.start.line,
+                character: selection.start.character,
+              },
+              end: {
+                line: selection.end.line,
+                character: selection.end.character,
+              },
             },
-          },
+          });
+          // }
         });
-        // }
       });
-    });
     return rangeInFiles;
   }
 
diff --git a/extension/src/debugPanel.ts b/extension/src/debugPanel.ts
index 487bbedf..dd24a8d8 100644
--- a/extension/src/debugPanel.ts
+++ b/extension/src/debugPanel.ts
@@ -6,78 +6,9 @@ import {
   openEditorAndRevealRange,
 } from "./util/vscode";
 import { RangeInFile } from "./client";
+import { setFocusedOnContinueInput } from "./commands";
 const WebSocket = require("ws");
 
-class StreamManager {
-  private _fullText: string = "";
-  private _insertionPoint: vscode.Position | undefined;
-
-  private _addToEditor(update: string) {
-    let editor =
-      vscode.window.activeTextEditor || vscode.window.visibleTextEditors[0];
-
-    if (typeof this._insertionPoint === "undefined") {
-      if (editor?.selection.isEmpty) {
-        this._insertionPoint = editor?.selection.active;
-      } else {
-        this._insertionPoint = editor?.selection.end;
-      }
-    }
-    editor?.edit((editBuilder) => {
-      if (this._insertionPoint) {
-        editBuilder.insert(this._insertionPoint, update);
-        this._insertionPoint = this._insertionPoint.translate(
-          Array.from(update.matchAll(/\n/g)).length,
-          update.length
-        );
-      }
-    });
-  }
-
-  public closeStream() {
-    this._fullText = "";
-    this._insertionPoint = undefined;
-    this._codeBlockStatus = "closed";
-    this._pendingBackticks = 0;
-  }
-
-  private _codeBlockStatus: "open" | "closed" | "language-descriptor" =
-    "closed";
-  private _pendingBackticks: number = 0;
-  public onStreamUpdate(update: string) {
-    let textToInsert = "";
-    for (let i = 0; i < update.length; i++) {
-      switch (this._codeBlockStatus) {
-        case "closed":
-          if (update[i] === "`" && this._fullText.endsWith("``")) {
-            this._codeBlockStatus = "language-descriptor";
-          }
-          break;
-        case "language-descriptor":
-          if (update[i] === " " || update[i] === "\n") {
-            this._codeBlockStatus = "open";
-          }
-          break;
-        case "open":
-          if (update[i] === "`") {
-            if (this._fullText.endsWith("``")) {
-              this._codeBlockStatus = "closed";
-              this._pendingBackticks = 0;
-            } else {
-              this._pendingBackticks += 1;
-            }
-          } else {
-            textToInsert += "`".repeat(this._pendingBackticks) + update[i];
-            this._pendingBackticks = 0;
-          }
-          break;
-      }
-      this._fullText += update[i];
-    }
-    this._addToEditor(textToInsert);
-  }
-}
-
 let websocketConnections: { [url: string]: WebsocketConnection | undefined } =
   {};
 
@@ -127,8 +58,6 @@ class WebsocketConnection {
   }
 }
 
-let streamManager = new StreamManager();
-
 export let debugPanelWebview: vscode.Webview | undefined;
 export function setupDebugPanel(
   panel: vscode.WebviewPanel | vscode.WebviewView,
@@ -147,10 +76,7 @@ export function setupDebugPanel(
     .toString();
 
   const isProduction = true; // context?.extensionMode === vscode.ExtensionMode.Development;
-  if (!isProduction) {
-    scriptUri = "http://localhost:5173/src/main.tsx";
-    styleMainUri = "http://localhost:5173/src/main.css";
-  } else {
+  if (isProduction) {
     scriptUri = debugPanelWebview
       .asWebviewUri(
         vscode.Uri.joinPath(extensionUri, "react-app/dist/assets/index.js")
@@ -161,6 +87,9 @@ export function setupDebugPanel(
         vscode.Uri.joinPath(extensionUri, "react-app/dist/assets/index.css")
       )
       .toString();
+  } else {
+    scriptUri = "http://localhost:5173/src/main.tsx";
+    styleMainUri = "http://localhost:5173/src/main.css";
   }
 
   panel.webview.options = {
@@ -175,11 +104,11 @@ export function setupDebugPanel(
       return;
     }
 
-    let rangeInFile: RangeInFile = {
+    const rangeInFile: RangeInFile = {
       range: e.selections[0],
       filepath: e.textEditor.document.fileName,
     };
-    let filesystem = {
+    const filesystem = {
       [rangeInFile.filepath]: e.textEditor.document.getText(),
     };
     panel.webview.postMessage({
@@ -217,13 +146,19 @@ export function setupDebugPanel(
           url,
         });
       };
-      const connection = new WebsocketConnection(
-        url,
-        onMessage,
-        onOpen,
-        onClose
-      );
-      websocketConnections[url] = connection;
+      try {
+        const connection = new WebsocketConnection(
+          url,
+          onMessage,
+          onOpen,
+          onClose
+        );
+        websocketConnections[url] = connection;
+        resolve(null);
+      } catch (e) {
+        console.log("Caught it!: ", e);
+        reject(e);
+      }
     });
   }
 
@@ -292,13 +227,8 @@ export function setupDebugPanel(
         openEditorAndRevealRange(data.path, undefined, vscode.ViewColumn.One);
         break;
       }
-      case "streamUpdate": {
-        // Write code at the position of the cursor
-        streamManager.onStreamUpdate(data.update);
-        break;
-      }
-      case "closeStream": {
-        streamManager.closeStream();
+      case "blurContinueInput": {
+        setFocusedOnContinueInput(false);
         break;
       }
       case "withProgress": {
diff --git a/extension/src/diffs.ts b/extension/src/diffs.ts
index b9ef8384..0bab326a 100644
--- a/extension/src/diffs.ts
+++ b/extension/src/diffs.ts
@@ -2,22 +2,31 @@ import * as os from "os";
 import * as path from "path";
 import * as fs from "fs";
 import * as vscode from "vscode";
-import { ideProtocolClient } from "./activation/activate";
+import { extensionContext, ideProtocolClient } from "./activation/activate";
+import { getMetaKeyLabel } from "./util/util";
+import { devDataPath } from "./activation/environmentSetup";
 
 interface DiffInfo {
   originalFilepath: string;
   newFilepath: string;
   editor?: vscode.TextEditor;
   step_index: number;
+  range: vscode.Range;
 }
 
-export const DIFF_DIRECTORY = path.join(os.homedir(), ".continue", "diffs");
+export const DIFF_DIRECTORY = path
+  .join(os.homedir(), ".continue", "diffs")
+  .replace(/^C:/, "c:");
 
 class DiffManager {
   // Create a temporary file in the global .continue directory which displays the updated version
   // Doing this because virtual files are read-only
   private diffs: Map<string, DiffInfo> = new Map();
 
+  diffAtNewFilepath(newFilepath: string): DiffInfo | undefined {
+    return this.diffs.get(newFilepath);
+  }
+
   private setupDirectory() {
     // Make sure the diff directory exists
     if (!fs.existsSync(DIFF_DIRECTORY)) {
@@ -35,6 +44,10 @@ class DiffManager {
     return filepath.replace(/\\/g, "_").replace(/\//g, "_");
   }
 
+  private getNewFilepath(originalFilepath: string): string {
+    return path.join(DIFF_DIRECTORY, this.escapeFilepath(originalFilepath));
+  }
+
   private openDiffEditor(
     originalFilepath: string,
     newFilepath: string
@@ -47,7 +60,7 @@ class DiffManager {
       return undefined;
     }
 
-    const rightUri = vscode.Uri.parse(newFilepath);
+    const rightUri = vscode.Uri.file(newFilepath);
     const leftUri = vscode.Uri.file(originalFilepath);
     const title = "Continue Diff";
     console.log(
@@ -70,9 +83,42 @@ class DiffManager {
       .getConfiguration("diffEditor", editor.document.uri)
       .update("codeLens", true, vscode.ConfigurationTarget.Global);
 
+    if (
+      extensionContext?.globalState.get<boolean>(
+        "continue.showDiffInfoMessage"
+      ) !== false
+    ) {
+      vscode.window
+        .showInformationMessage(
+          `Accept (${getMetaKeyLabel()}⇧↩) or reject (${getMetaKeyLabel()}⇧⌫) at the top of the file.`,
+          "Got it",
+          "Don't show again"
+        )
+        .then((selection) => {
+          if (selection === "Don't show again") {
+            // Get the global state
+            extensionContext?.globalState.update(
+              "continue.showDiffInfoMessage",
+              false
+            );
+          }
+        });
+    }
+
     return editor;
   }
 
+  private _findFirstDifferentLine(contentA: string, contentB: string): number {
+    const linesA = contentA.split("\n");
+    const linesB = contentB.split("\n");
+    for (let i = 0; i < linesA.length && i < linesB.length; i++) {
+      if (linesA[i] !== linesB[i]) {
+        return i;
+      }
+    }
+    return 0;
+  }
+
   writeDiff(
     originalFilepath: string,
     newContent: string,
@@ -81,18 +127,20 @@ class DiffManager {
     this.setupDirectory();
 
     // Create or update existing diff
-    const newFilepath = path.join(
-      DIFF_DIRECTORY,
-      this.escapeFilepath(originalFilepath)
-    );
+    const newFilepath = this.getNewFilepath(originalFilepath);
     fs.writeFileSync(newFilepath, newContent);
 
     // Open the diff editor if this is a new diff
     if (!this.diffs.has(newFilepath)) {
+      // Figure out the first line that is different
+      const oldContent = ideProtocolClient.readFile(originalFilepath);
+      const line = this._findFirstDifferentLine(oldContent, newContent);
+
       const diffInfo: DiffInfo = {
         originalFilepath,
         newFilepath,
         step_index,
+        range: new vscode.Range(line, 0, line + 1, 0),
       };
       this.diffs.set(newFilepath, diffInfo);
     }
@@ -104,6 +152,11 @@ class DiffManager {
       this.diffs.set(newFilepath, diffInfo);
     }
 
+    vscode.commands.executeCommand(
+      "workbench.action.files.revert",
+      vscode.Uri.file(newFilepath)
+    );
+
     return newFilepath;
   }
 
@@ -117,10 +170,38 @@ class DiffManager {
     fs.unlinkSync(diffInfo.newFilepath);
   }
 
+  private inferNewFilepath() {
+    const activeEditorPath =
+      vscode.window.activeTextEditor?.document.uri.fsPath;
+    if (activeEditorPath && path.dirname(activeEditorPath) === DIFF_DIRECTORY) {
+      return activeEditorPath;
+    }
+    const visibleEditors = vscode.window.visibleTextEditors.map(
+      (editor) => editor.document.uri.fsPath
+    );
+    for (const editorPath of visibleEditors) {
+      if (path.dirname(editorPath) === DIFF_DIRECTORY) {
+        for (const otherEditorPath of visibleEditors) {
+          if (
+            path.dirname(otherEditorPath) !== DIFF_DIRECTORY &&
+            this.getNewFilepath(otherEditorPath) === editorPath
+          ) {
+            return editorPath;
+          }
+        }
+      }
+    }
+
+    if (this.diffs.size === 1) {
+      return Array.from(this.diffs.keys())[0];
+    }
+    return undefined;
+  }
+
   acceptDiff(newFilepath?: string) {
-    // If no newFilepath is provided and there is only one in the dictionary, use that
-    if (!newFilepath && this.diffs.size === 1) {
-      newFilepath = Array.from(this.diffs.keys())[0];
+    // When coming from a keyboard shortcut, we have to infer the newFilepath from visible text editors
+    if (!newFilepath) {
+      newFilepath = this.inferNewFilepath();
     }
     if (!newFilepath) {
       console.log("No newFilepath provided to accept the diff");
@@ -132,20 +213,32 @@ class DiffManager {
       console.log("No corresponding diffInfo found for newFilepath");
       return;
     }
-    fs.writeFileSync(
-      diffInfo.originalFilepath,
-      fs.readFileSync(diffInfo.newFilepath)
-    );
-    this.cleanUpDiff(diffInfo);
+
+    // Save the right-side file, then copy over to original
+    vscode.workspace.textDocuments
+      .find((doc) => doc.uri.fsPath === newFilepath)
+      ?.save()
+      .then(() => {
+        fs.writeFileSync(
+          diffInfo.originalFilepath,
+          fs.readFileSync(diffInfo.newFilepath)
+        );
+        this.cleanUpDiff(diffInfo);
+      });
+
+    recordAcceptReject(true, diffInfo);
   }
 
   rejectDiff(newFilepath?: string) {
     // If no newFilepath is provided and there is only one in the dictionary, use that
-    if (!newFilepath && this.diffs.size === 1) {
-      newFilepath = Array.from(this.diffs.keys())[0];
+    if (!newFilepath) {
+      newFilepath = this.inferNewFilepath();
     }
     if (!newFilepath) {
-      console.log("No newFilepath provided to reject the diff");
+      console.log(
+        "No newFilepath provided to reject the diff, diffs.size was",
+        this.diffs.size
+      );
       return;
     }
     const diffInfo = this.diffs.get(newFilepath);
@@ -157,12 +250,56 @@ class DiffManager {
     // Stop the step at step_index in case it is still streaming
     ideProtocolClient.deleteAtIndex(diffInfo.step_index);
 
-    this.cleanUpDiff(diffInfo);
+    vscode.workspace.textDocuments
+      .find((doc) => doc.uri.fsPath === newFilepath)
+      ?.save()
+      .then(() => {
+        this.cleanUpDiff(diffInfo);
+      });
+
+    recordAcceptReject(false, diffInfo);
   }
 }
 
 export const diffManager = new DiffManager();
 
+function recordAcceptReject(accepted: boolean, diffInfo: DiffInfo) {
+  const collectOn = vscode.workspace
+    .getConfiguration("continue")
+    .get<boolean>("dataSwitch");
+
+  if (collectOn) {
+    const devDataDir = devDataPath();
+    const suggestionsPath = path.join(devDataDir, "suggestions.json");
+
+    // Initialize suggestions list
+    let suggestions = [];
+
+    // Check if suggestions.json exists
+    if (fs.existsSync(suggestionsPath)) {
+      const rawData = fs.readFileSync(suggestionsPath, "utf-8");
+      suggestions = JSON.parse(rawData);
+    }
+
+    // Add the new suggestion to the list
+    suggestions.push({
+      accepted,
+      timestamp: Date.now(),
+      suggestion: diffInfo.originalFilepath,
+    });
+
+    // Send the suggestion to the server
+    ideProtocolClient.sendAcceptRejectSuggestion(accepted);
+
+    // Write the updated suggestions back to the file
+    fs.writeFileSync(
+      suggestionsPath,
+      JSON.stringify(suggestions, null, 4),
+      "utf-8"
+    );
+  }
+}
+
 export async function acceptDiffCommand(newFilepath?: string) {
   diffManager.acceptDiff(newFilepath);
   ideProtocolClient.sendAcceptRejectDiff(true);
diff --git a/extension/src/extension.ts b/extension/src/extension.ts
index de8f55e3..6959ec05 100644
--- a/extension/src/extension.ts
+++ b/extension/src/extension.ts
@@ -3,17 +3,17 @@
  */
 
 import * as vscode from "vscode";
-import {
-  isPythonEnvSetup,
-  startContinuePythonServer,
-} from "./activation/environmentSetup";
 
-async function dynamicImportAndActivate(
-  context: vscode.ExtensionContext,
-  showTutorial: boolean
-) {
+async function dynamicImportAndActivate(context: vscode.ExtensionContext) {
   const { activateExtension } = await import("./activation/activate");
-  await activateExtension(context, showTutorial);
+  try {
+    await activateExtension(context);
+  } catch (e) {
+    console.log("Error activating extension: ", e);
+    vscode.window.showInformationMessage(
+      "Error activating the Continue extension."
+    );
+  }
 }
 
 export function activate(context: vscode.ExtensionContext) {
@@ -25,7 +25,7 @@ export function activate(context: vscode.ExtensionContext) {
       cancellable: false,
     },
     async () => {
-      dynamicImportAndActivate(context, true);
+      dynamicImportAndActivate(context);
     }
   );
 }
diff --git a/extension/src/lang-server/codeActions.ts b/extension/src/lang-server/codeActions.ts
new file mode 100644
index 00000000..07cf5f4e
--- /dev/null
+++ b/extension/src/lang-server/codeActions.ts
@@ -0,0 +1,53 @@
+import * as vscode from "vscode";
+
+class ContinueQuickFixProvider implements vscode.CodeActionProvider {
+  public static readonly providedCodeActionKinds = [
+    vscode.CodeActionKind.QuickFix,
+  ];
+
+  provideCodeActions(
+    document: vscode.TextDocument,
+    range: vscode.Range | vscode.Selection,
+    context: vscode.CodeActionContext,
+    token: vscode.CancellationToken
+  ): vscode.ProviderResult<(vscode.Command | vscode.CodeAction)[]> {
+    if (context.diagnostics.length === 0) {
+      return [];
+    }
+
+    const createQuickFix = (edit: boolean) => {
+      const diagnostic = context.diagnostics[0];
+      const quickFix = new vscode.CodeAction(
+        edit ? "Fix with Continue" : "Ask Continue",
+        vscode.CodeActionKind.QuickFix
+      );
+      quickFix.isPreferred = false;
+      const surroundingRange = new vscode.Range(
+        range.start.translate(-3, 0),
+        range.end.translate(3, 0)
+      );
+      quickFix.command = {
+        command: "continue.quickFix",
+        title: "Continue Quick Fix",
+        arguments: [
+          diagnostic.message,
+          document.getText(surroundingRange),
+          edit,
+        ],
+      };
+      return quickFix;
+    };
+    return [createQuickFix(true), createQuickFix(false)];
+  }
+}
+
+export default function registerQuickFixProvider() {
+  // In your extension's activate function:
+  vscode.languages.registerCodeActionsProvider(
+    { language: "*" },
+    new ContinueQuickFixProvider(),
+    {
+      providedCodeActionKinds: ContinueQuickFixProvider.providedCodeActionKinds,
+    }
+  );
+}
diff --git a/extension/src/lang-server/codeLens.ts b/extension/src/lang-server/codeLens.ts
index 79126eaa..ba80e557 100644
--- a/extension/src/lang-server/codeLens.ts
+++ b/extension/src/lang-server/codeLens.ts
@@ -2,11 +2,12 @@ import * as vscode from "vscode";
 import { editorToSuggestions, editorSuggestionsLocked } from "../suggestions";
 import * as path from "path";
 import * as os from "os";
-import { DIFF_DIRECTORY } from "../diffs";
+import { DIFF_DIRECTORY, diffManager } from "../diffs";
+import { getMetaKeyLabel } from "../util/util";
 class SuggestionsCodeLensProvider implements vscode.CodeLensProvider {
   public provideCodeLenses(
     document: vscode.TextDocument,
-    token: vscode.CancellationToken
+    _: vscode.CancellationToken
   ): vscode.CodeLens[] | Thenable<vscode.CodeLens[]> {
     const suggestions = editorToSuggestions.get(document.uri.toString());
     if (!suggestions) {
@@ -35,7 +36,7 @@ class SuggestionsCodeLensProvider implements vscode.CodeLensProvider {
       if (codeLenses.length === 2) {
         codeLenses.push(
           new vscode.CodeLens(range, {
-            title: "(⌘⇧↩/⌘⇧⌫ to accept/reject all)",
+            title: `(${getMetaKeyLabel()}⇧↩/${getMetaKeyLabel()}⇧⌫ to accept/reject all)`,
             command: "",
           })
         );
@@ -44,40 +45,28 @@ class SuggestionsCodeLensProvider implements vscode.CodeLensProvider {
 
     return codeLenses;
   }
-
-  onDidChangeCodeLenses?: vscode.Event<void> | undefined;
-
-  constructor(emitter?: vscode.EventEmitter<void>) {
-    if (emitter) {
-      this.onDidChangeCodeLenses = emitter.event;
-      this.onDidChangeCodeLenses(() => {
-        if (vscode.window.activeTextEditor) {
-          this.provideCodeLenses(
-            vscode.window.activeTextEditor.document,
-            new vscode.CancellationTokenSource().token
-          );
-        }
-      });
-    }
-  }
 }
 
 class DiffViewerCodeLensProvider implements vscode.CodeLensProvider {
   public provideCodeLenses(
     document: vscode.TextDocument,
-    token: vscode.CancellationToken
+    _: vscode.CancellationToken
   ): vscode.CodeLens[] | Thenable<vscode.CodeLens[]> {
     if (path.dirname(document.uri.fsPath) === DIFF_DIRECTORY) {
       const codeLenses: vscode.CodeLens[] = [];
-      const range = new vscode.Range(0, 0, 1, 0);
+      let range = new vscode.Range(0, 0, 1, 0);
+      const diffInfo = diffManager.diffAtNewFilepath(document.uri.fsPath);
+      if (diffInfo) {
+        range = diffInfo.range;
+      }
       codeLenses.push(
         new vscode.CodeLens(range, {
-          title: "Accept ✅ (⌘⇧↩)",
+          title: `Accept All ✅ (${getMetaKeyLabel()}⇧↩)`,
           command: "continue.acceptDiff",
           arguments: [document.uri.fsPath],
         }),
         new vscode.CodeLens(range, {
-          title: "Reject ❌ (⌘⇧⌫)",
+          title: `Reject All ❌ (${getMetaKeyLabel()}⇧⌫)`,
           command: "continue.rejectDiff",
           arguments: [document.uri.fsPath],
         })
@@ -87,22 +76,6 @@ class DiffViewerCodeLensProvider implements vscode.CodeLensProvider {
       return [];
     }
   }
-
-  onDidChangeCodeLenses?: vscode.Event<void> | undefined;
-
-  constructor(emitter?: vscode.EventEmitter<void>) {
-    if (emitter) {
-      this.onDidChangeCodeLenses = emitter.event;
-      this.onDidChangeCodeLenses(() => {
-        if (vscode.window.activeTextEditor) {
-          this.provideCodeLenses(
-            vscode.window.activeTextEditor.document,
-            new vscode.CancellationTokenSource().token
-          );
-        }
-      });
-    }
-  }
 }
 
 let diffsCodeLensDisposable: vscode.Disposable | undefined = undefined;
diff --git a/extension/src/suggestions.ts b/extension/src/suggestions.ts
index 6e5a444f..c2373223 100644
--- a/extension/src/suggestions.ts
+++ b/extension/src/suggestions.ts
@@ -1,9 +1,7 @@
 import * as vscode from "vscode";
 import { sendTelemetryEvent, TelemetryEvent } from "./telemetry";
 import { openEditorAndRevealRange } from "./util/vscode";
-import { translate, readFileAtRange } from "./util/vscode";
-import * as fs from "fs";
-import * as path from "path";
+import { translate } from "./util/vscode";
 import { registerAllCodeLensProviders } from "./lang-server/codeLens";
 import { extensionContext, ideProtocolClient } from "./activation/activate";
 
@@ -214,62 +212,6 @@ function selectSuggestion(
         : suggestion.newRange;
   }
 
-  let workspaceDir = vscode.workspace.workspaceFolders
-    ? vscode.workspace.workspaceFolders[0]?.uri.fsPath
-    : undefined;
-
-  let collectOn = vscode.workspace
-    .getConfiguration("continue")
-    .get<boolean>("dataSwitch");
-
-  if (workspaceDir && collectOn) {
-    let continueDir = path.join(workspaceDir, ".continue");
-
-    // Check if .continue directory doesn't exists
-    if (!fs.existsSync(continueDir)) {
-      fs.mkdirSync(continueDir);
-    }
-
-    let suggestionsPath = path.join(continueDir, "suggestions.json");
-
-    // Initialize suggestions list
-    let suggestions = [];
-
-    // Check if suggestions.json exists
-    if (fs.existsSync(suggestionsPath)) {
-      let rawData = fs.readFileSync(suggestionsPath, "utf-8");
-      suggestions = JSON.parse(rawData);
-    }
-
-    const accepted =
-      accept === "new" || (accept === "selected" && suggestion.newSelected);
-    suggestions.push({
-      accepted,
-      timestamp: Date.now(),
-      suggestion: suggestion.newContent,
-    });
-    ideProtocolClient.sendAcceptRejectSuggestion(accepted);
-
-    // Write the updated suggestions back to the file
-    fs.writeFileSync(
-      suggestionsPath,
-      JSON.stringify(suggestions, null, 4),
-      "utf-8"
-    );
-
-    // If it's not already there, add .continue to .gitignore
-    const gitignorePath = path.join(workspaceDir, ".gitignore");
-    if (fs.existsSync(gitignorePath)) {
-      const gitignoreData = fs.readFileSync(gitignorePath, "utf-8");
-      const gitIgnoreLines = gitignoreData.split("\n");
-      if (!gitIgnoreLines.includes(".continue")) {
-        fs.appendFileSync(gitignorePath, "\n.continue\n");
-      }
-    } else {
-      fs.writeFileSync(gitignorePath, ".continue\n");
-    }
-  }
-
   rangeToDelete = new vscode.Range(
     rangeToDelete.start,
     new vscode.Position(rangeToDelete.end.line, 0)
diff --git a/extension/src/util/messenger.ts b/extension/src/util/messenger.ts
index b1df161b..3044898e 100644
--- a/extension/src/util/messenger.ts
+++ b/extension/src/util/messenger.ts
@@ -16,6 +16,8 @@ export abstract class Messenger {
 
   abstract onClose(callback: () => void): void;
 
+  abstract onError(callback: () => void): void;
+
   abstract sendAndReceive(messageType: string, data: any): Promise<any>;
 }
 
@@ -26,6 +28,7 @@ export class WebsocketMessenger extends Messenger {
   } = {};
   private onOpenListeners: (() => void)[] = [];
   private onCloseListeners: (() => void)[] = [];
+  private onErrorListeners: (() => void)[] = [];
   private serverUrl: string;
 
   _newWebsocket(): WebSocket {
@@ -43,6 +46,9 @@ export class WebsocketMessenger extends Messenger {
     for (const listener of this.onCloseListeners) {
       this.onClose(listener);
     }
+    for (const listener of this.onErrorListeners) {
+      this.onError(listener);
+    }
     for (const messageType in this.onMessageListeners) {
       for (const listener of this.onMessageListeners[messageType]) {
         this.onMessageType(messageType, listener);
@@ -151,4 +157,8 @@ export class WebsocketMessenger extends Messenger {
   onClose(callback: () => void): void {
     this.websocket.addEventListener("close", callback);
   }
+
+  onError(callback: () => void): void {
+    this.websocket.addEventListener("error", callback);
+  }
 }
diff --git a/extension/src/util/util.ts b/extension/src/util/util.ts
index d33593e1..dfc10c90 100644
--- a/extension/src/util/util.ts
+++ b/extension/src/util/util.ts
@@ -1,5 +1,6 @@
 import { RangeInFile, SerializedDebugContext } from "../client";
 import * as fs from "fs";
+const os = require("os");
 
 function charIsEscapedAtIndex(index: number, str: string): boolean {
   if (index === 0) return false;
@@ -113,3 +114,31 @@ export function debounced(delay: number, fn: Function) {
     }, delay);
   };
 }
+
+type Platform = "mac" | "linux" | "windows" | "unknown";
+
+function getPlatform(): Platform {
+  const platform = os.platform();
+  if (platform === "darwin") {
+    return "mac";
+  } else if (platform === "linux") {
+    return "linux";
+  } else if (platform === "win32") {
+    return "windows";
+  } else {
+    return "unknown";
+  }
+}
+
+export function getMetaKeyLabel() {
+  const platform = getPlatform();
+  switch (platform) {
+    case "mac":
+      return "⌘";
+    case "linux":
+    case "windows":
+      return "^";
+    default:
+      return "⌘";
+  }
+}