Merge pull request #243 from continuedev/reconnect-persist-state

Reconnect persist state

18 files changed, 196 insertions, 157 deletions

diff --git a/continuedev/src/continuedev/core/autopilot.py b/continuedev/src/continuedev/core/autopilot.py
index 1b074435..e1c8a076 100644
--- a/continuedev/src/continuedev/core/autopilot.py
+++ b/continuedev/src/continuedev/core/autopilot.py
@@ -1,13 +1,13 @@
 from functools import cached_property
 import traceback
 import time
-from typing import Any, Callable, Coroutine, Dict, List
+from typing import Any, Callable, Coroutine, Dict, List, Union
 import os
 from aiohttp import ClientPayloadError
+from pydantic import root_validator
 
 from ..models.filesystem import RangeInFileWithContents
 from ..models.filesystem_edit import FileEditWithFullContents
-from ..libs.llm import LLM
 from .observation import Observation, InternalErrorObservation
 from ..server.ide_protocol import AbstractIdeProtocolServer
 from ..libs.util.queue import AsyncSubscriptionQueue
@@ -16,7 +16,6 @@ from .main import Context, ContinueCustomException, HighlightedRangeContext, Pol
 from ..steps.core.core import ReversibleStep, ManualEditStep, UserInputStep
 from ..libs.util.telemetry import capture_event
 from .sdk import ContinueSDK
-import asyncio
 from ..libs.util.step_name_to_steps import get_step_from_name
 from ..libs.util.traceback_parsers import get_python_traceback, get_javascript_traceback
 from openai import error as openai_errors
@@ -46,6 +45,7 @@ class Autopilot(ContinueBaseModel):
     ide: AbstractIdeProtocolServer
     history: History = History.from_empty()
     context: Context = Context()
+    full_state: Union[FullState, None] = None
     _on_update_callbacks: List[Callable[[FullState], None]] = []
 
     _active: bool = False
@@ -63,8 +63,15 @@ class Autopilot(ContinueBaseModel):
         arbitrary_types_allowed = True
         keep_untouched = (cached_property,)
 
+    @root_validator(pre=True)
+    def fill_in_values(cls, values):
+        full_state: FullState = values.get('full_state')
+        if full_state is not None:
+            values['history'] = full_state.history
+        return values
+
     def get_full_state(self) -> FullState:
-        return FullState(
+        full_state = FullState(
             history=self.history,
             active=self._active,
             user_input_queue=self._main_user_input_queue,
@@ -73,6 +80,8 @@ class Autopilot(ContinueBaseModel):
             slash_commands=self.get_available_slash_commands(),
             adding_highlighted_code=self._adding_highlighted_code,
         )
+        self.full_state = full_state
+        return full_state
 
     def get_available_slash_commands(self) -> List[Dict]:
         custom_commands = list(map(lambda x: {
diff --git a/continuedev/src/continuedev/libs/constants/main.py b/continuedev/src/continuedev/libs/constants/main.py
new file mode 100644
index 00000000..96eb6e69
--- /dev/null
+++ b/continuedev/src/continuedev/libs/constants/main.py
@@ -0,0 +1,6 @@
+## PATHS ##
+
+CONTINUE_GLOBAL_FOLDER = ".continue"
+CONTINUE_SESSIONS_FOLDER = "sessions"
+CONTINUE_SERVER_FOLDER = "server"
+
diff --git a/continuedev/src/continuedev/libs/util/paths.py b/continuedev/src/continuedev/libs/util/paths.py
new file mode 100644
index 00000000..fddef887
--- /dev/null
+++ b/continuedev/src/continuedev/libs/util/paths.py
@@ -0,0 +1,17 @@
+import os
+
+from ..constants.main import CONTINUE_SESSIONS_FOLDER, CONTINUE_GLOBAL_FOLDER, CONTINUE_SERVER_FOLDER
+
+def getGlobalFolderPath():
+    return os.path.join(os.path.expanduser("~"), CONTINUE_GLOBAL_FOLDER)
+
+
+
+def getSessionsFolderPath():
+    return os.path.join(getGlobalFolderPath(), CONTINUE_SESSIONS_FOLDER)
+
+def getServerFolderPath():
+    return os.path.join(getGlobalFolderPath(), CONTINUE_SERVER_FOLDER)
+
+def getSessionFilePath(session_id: str):
+    return os.path.join(getSessionsFolderPath(), f"{session_id}.json")
\ No newline at end of file
diff --git a/continuedev/src/continuedev/server/gui.py b/continuedev/src/continuedev/server/gui.py
index 21089f30..238273b2 100644
--- a/continuedev/src/continuedev/server/gui.py
+++ b/continuedev/src/continuedev/server/gui.py
@@ -31,12 +31,12 @@ class AppStatus:
 Server.handle_exit = AppStatus.handle_exit
 
 
-def session(x_continue_session_id: str = Header("anonymous")) -> Session:
-    return session_manager.get_session(x_continue_session_id)
+async def session(x_continue_session_id: str = Header("anonymous")) -> Session:
+    return await session_manager.get_session(x_continue_session_id)
 
 
-def websocket_session(session_id: str) -> Session:
-    return session_manager.get_session(session_id)
+async def websocket_session(session_id: str) -> Session:
+    return await session_manager.get_session(session_id)
 
 
 T = TypeVar("T", bound=BaseModel)
@@ -174,7 +174,7 @@ async def websocket_endpoint(websocket: WebSocket, session: Session = Depends(we
         protocol.websocket = websocket
 
         # Update any history that may have happened before connection
-        # await protocol.send_state_update()
+        await protocol.send_state_update()
 
         while AppStatus.should_exit is False:
             message = await websocket.receive_text()
@@ -199,4 +199,6 @@ async def websocket_endpoint(websocket: WebSocket, session: Session = Depends(we
         print("Closing gui websocket")
         if websocket.client_state != WebSocketState.DISCONNECTED:
             await websocket.close()
+
+        session_manager.persist_session(session.session_id)
         session_manager.remove_session(session.session_id)
diff --git a/continuedev/src/continuedev/server/ide.py b/continuedev/src/continuedev/server/ide.py
index 4645b49e..12a21f19 100644
--- a/continuedev/src/continuedev/server/ide.py
+++ b/continuedev/src/continuedev/server/ide.py
@@ -52,9 +52,11 @@ class FileEditsUpdate(BaseModel):
 class OpenFilesResponse(BaseModel):
     openFiles: List[str]
 
+
 class VisibleFilesResponse(BaseModel):
     visibleFiles: List[str]
 
+
 class HighlightedCodeResponse(BaseModel):
     highlightedCode: List[RangeInFile]
 
@@ -115,6 +117,7 @@ class IdeProtocolServer(AbstractIdeProtocolServer):
     websocket: WebSocket
     session_manager: SessionManager
     sub_queue: AsyncSubscriptionQueue = AsyncSubscriptionQueue()
+    session_id: Union[str, None] = None
 
     def __init__(self, session_manager: SessionManager, websocket: WebSocket):
         self.websocket = websocket
@@ -132,8 +135,6 @@ class IdeProtocolServer(AbstractIdeProtocolServer):
                 continue
             message_type = message["messageType"]
             data = message["data"]
-            # if message_type == "openGUI":
-            #     await self.openGUI()
             if message_type == "workspaceDirectory":
                 self.workspace_directory = data["workspaceDirectory"]
                 break
@@ -158,8 +159,8 @@ class IdeProtocolServer(AbstractIdeProtocolServer):
         return resp_model.parse_obj(resp)
 
     async def handle_json(self, message_type: str, data: Any):
-        if message_type == "openGUI":
-            await self.openGUI()
+        if message_type == "getSessionId":
+            await self.getSessionId()
         elif message_type == "setFileOpen":
             await self.setFileOpen(data["filepath"], data["open"])
         elif message_type == "setSuggestionsLocked":
@@ -217,9 +218,10 @@ class IdeProtocolServer(AbstractIdeProtocolServer):
             "locked": locked
         })
 
-    async def openGUI(self):
-        session_id = self.session_manager.new_session(self)
-        await self._send_json("openGUI", {
+    async def getSessionId(self):
+        session_id = self.session_manager.new_session(
+            self, self.session_id).session_id
+        await self._send_json("getSessionId", {
             "sessionId": session_id
         })
 
@@ -304,7 +306,7 @@ class IdeProtocolServer(AbstractIdeProtocolServer):
     async def getOpenFiles(self) -> List[str]:
         resp = await self._send_and_receive_json({}, OpenFilesResponse, "openFiles")
         return resp.openFiles
-    
+
     async def getVisibleFiles(self) -> List[str]:
         resp = await self._send_and_receive_json({}, VisibleFilesResponse, "visibleFiles")
         return resp.visibleFiles
@@ -416,7 +418,7 @@ class IdeProtocolServer(AbstractIdeProtocolServer):
 
 
 @router.websocket("/ws")
-async def websocket_endpoint(websocket: WebSocket):
+async def websocket_endpoint(websocket: WebSocket, session_id: str = None):
     try:
         await websocket.accept()
         print("Accepted websocket connection from, ", websocket.client)
@@ -434,6 +436,9 @@ async def websocket_endpoint(websocket: WebSocket):
                 ideProtocolServer.handle_json(message_type, data))
 
         ideProtocolServer = IdeProtocolServer(session_manager, websocket)
+        ideProtocolServer.session_id = session_id
+        if session_id is not None:
+            session_manager.registered_ides[session_id] = ideProtocolServer
         other_msgs = await ideProtocolServer.initialize()
 
         for other_msg in other_msgs:
@@ -454,3 +459,5 @@ async def websocket_endpoint(websocket: WebSocket):
     finally:
         if websocket.client_state != WebSocketState.DISCONNECTED:
             await websocket.close()
+
+        session_manager.registered_ides.pop(ideProtocolServer.session_id)
diff --git a/continuedev/src/continuedev/server/ide_protocol.py b/continuedev/src/continuedev/server/ide_protocol.py
index 2783dc61..2f78cf0e 100644
--- a/continuedev/src/continuedev/server/ide_protocol.py
+++ b/continuedev/src/continuedev/server/ide_protocol.py
@@ -1,5 +1,6 @@
-from typing import Any, List
+from typing import Any, List, Union
 from abc import ABC, abstractmethod, abstractproperty
+from fastapi import WebSocket
 
 from ..models.main import Traceback
 from ..models.filesystem_edit import FileEdit, FileSystemEdit, EditDiff
@@ -7,6 +8,9 @@ from ..models.filesystem import RangeInFile, RangeInFileWithContents
 
 
 class AbstractIdeProtocolServer(ABC):
+    websocket: WebSocket
+    session_id: Union[str, None]
+
     @abstractmethod
     async def handle_json(self, data: Any):
         """Handle a json message"""
@@ -24,8 +28,8 @@ class AbstractIdeProtocolServer(ABC):
         """Set whether suggestions are locked"""
 
     @abstractmethod
-    async def openGUI(self):
-        """Open a GUI"""
+    async def getSessionId(self):
+        """Get a new session ID"""
 
     @abstractmethod
     async def showSuggestionsAndWait(self, suggestions: List[FileEdit]) -> bool:
diff --git a/continuedev/src/continuedev/server/main.py b/continuedev/src/continuedev/server/main.py
index f4d82903..aa093853 100644
--- a/continuedev/src/continuedev/server/main.py
+++ b/continuedev/src/continuedev/server/main.py
@@ -4,7 +4,8 @@ from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from .ide import router as ide_router
 from .gui import router as gui_router
-import logging
+from .session_manager import session_manager
+import atexit
 import uvicorn
 import argparse
 
@@ -44,5 +45,16 @@ def run_server():
     uvicorn.run(app, host="0.0.0.0", port=args.port)
 
 
+def cleanup():
+    print("Cleaning up sessions")
+    for session_id in session_manager.sessions:
+        session_manager.persist_session(session_id)
+
+
+atexit.register(cleanup)
 if __name__ == "__main__":
-    run_server()
+    try:
+        run_server()
+    except Exception as e:
+        cleanup()
+        raise e
diff --git a/continuedev/src/continuedev/server/session_manager.py b/continuedev/src/continuedev/server/session_manager.py
index 7147dcfa..fb8ac386 100644
--- a/continuedev/src/continuedev/server/session_manager.py
+++ b/continuedev/src/continuedev/server/session_manager.py
@@ -1,9 +1,12 @@
-from asyncio import BaseEventLoop
+import os
 from fastapi import WebSocket
 from typing import Any, Dict, List, Union
 from uuid import uuid4
+import json
 
+from ..libs.util.paths import getSessionFilePath, getSessionsFolderPath
 from ..models.filesystem_edit import FileEditWithFullContents
+from ..libs.constants.main import CONTINUE_SESSIONS_FOLDER
 from ..core.policy import DemoPolicy
 from ..core.main import FullState
 from ..core.autopilot import Autopilot
@@ -39,17 +42,35 @@ class DemoAutopilot(Autopilot):
 
 class SessionManager:
     sessions: Dict[str, Session] = {}
+    # Mapping of session_id to IDE, where the IDE is still alive
+    registered_ides: Dict[str, AbstractIdeProtocolServer] = {}
 
-    def get_session(self, session_id: str) -> Session:
+    async def get_session(self, session_id: str) -> Session:
         if session_id not in self.sessions:
+            # Check then whether it is persisted by listing all files in the sessions folder
+            # And only if the IDE is still alive
+            sessions_folder = getSessionsFolderPath()
+            session_files = os.listdir(sessions_folder)
+            if f"{session_id}.json" in session_files and session_id in self.registered_ides:
+                if self.registered_ides[session_id].session_id is not None:
+                    return self.new_session(self.registered_ides[session_id], session_id=session_id)
+
             raise KeyError("Session ID not recognized", session_id)
         return self.sessions[session_id]
 
-    def new_session(self, ide: AbstractIdeProtocolServer) -> str:
-        autopilot = DemoAutopilot(policy=DemoPolicy(), ide=ide)
-        session_id = str(uuid4())
+    def new_session(self, ide: AbstractIdeProtocolServer, session_id: Union[str, None] = None) -> Session:
+        full_state = None
+        if session_id is not None and os.path.exists(getSessionFilePath(session_id)):
+            with open(getSessionFilePath(session_id), "r") as f:
+                full_state = FullState(**json.load(f))
+
+        autopilot = DemoAutopilot(
+            policy=DemoPolicy(), ide=ide, full_state=full_state)
+        session_id = session_id or str(uuid4())
+        ide.session_id = session_id
         session = Session(session_id=session_id, autopilot=autopilot)
         self.sessions[session_id] = session
+        self.registered_ides[session_id] = ide
 
         async def on_update(state: FullState):
             await session_manager.send_ws_data(session_id, "state_update", {
@@ -58,11 +79,19 @@ class SessionManager:
 
         autopilot.on_update(on_update)
         create_async_task(autopilot.run_policy())
-        return session_id
+        return session
 
     def remove_session(self, session_id: str):
         del self.sessions[session_id]
 
+    def persist_session(self, session_id: str):
+        """Save the session's FullState as a json file"""
+        full_state = self.sessions[session_id].autopilot.get_full_state()
+        if not os.path.exists(getSessionsFolderPath()):
+            os.mkdir(getSessionsFolderPath())
+        with open(getSessionFilePath(session_id), "w") as f:
+            json.dump(full_state.dict(), f)
+
     def register_websocket(self, session_id: str, ws: WebSocket):
         self.sessions[session_id].ws = ws
         print("Registered websocket for session", session_id)
diff --git a/docs/docs/concepts/ide.md b/docs/docs/concepts/ide.md
index dc7b9e23..bd31481b 100644
--- a/docs/docs/concepts/ide.md
+++ b/docs/docs/concepts/ide.md
@@ -41,9 +41,9 @@ Get the workspace directory
 

 Set whether a file is open

 

-### openGUI

+### getSessionId

 

-Open a gui

+Get a new session ID

 

 ### showSuggestionsAndWait

 

diff --git a/extension/package-lock.json b/extension/package-lock.json
index 01e5bc1a..ea853fae 100644
--- a/extension/package-lock.json
+++ b/extension/package-lock.json
@@ -1,12 +1,12 @@
 {
   "name": "continue",
-  "version": "0.0.156",
+  "version": "0.0.157",
   "lockfileVersion": 2,
   "requires": true,
   "packages": {
     "": {
       "name": "continue",
-      "version": "0.0.156",
+      "version": "0.0.157",
       "license": "Apache-2.0",
       "dependencies": {
         "@electron/rebuild": "^3.2.10",
diff --git a/extension/package.json b/extension/package.json
index f3402abf..b3ff0b52 100644
--- a/extension/package.json
+++ b/extension/package.json
@@ -14,7 +14,7 @@
   "displayName": "Continue",
   "pricing": "Free",
   "description": "The open-source coding autopilot",
-  "version": "0.0.156",
+  "version": "0.0.157",
   "publisher": "Continue",
   "engines": {
     "vscode": "^1.67.0"
diff --git a/extension/react-app/src/components/ComboBox.tsx b/extension/react-app/src/components/ComboBox.tsx
index ac994b0a..7d6541c7 100644
--- a/extension/react-app/src/components/ComboBox.tsx
+++ b/extension/react-app/src/components/ComboBox.tsx
@@ -331,7 +331,6 @@ const ComboBox = React.forwardRef((props: ComboBoxProps, ref) => {
               ) {
                 (event.nativeEvent as any).preventDownshiftDefault = true;
               } else if (event.key === "ArrowUp") {
-                console.log("OWJFOIJO");
                 if (positionInHistory == 0) return;
                 else if (
                   positionInHistory == history.length &&
diff --git a/extension/react-app/src/hooks/messenger.ts b/extension/react-app/src/hooks/messenger.ts
index e2a0bab8..00ce1fbb 100644
--- a/extension/react-app/src/hooks/messenger.ts
+++ b/extension/react-app/src/hooks/messenger.ts
@@ -1,6 +1,3 @@
-// console.log("Websocket import");
-// const WebSocket = require("ws");
-
 export abstract class Messenger {
   abstract send(messageType: string, data: object): void;
 
@@ -28,13 +25,6 @@ export class WebsocketMessenger extends Messenger {
   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);
diff --git a/extension/react-app/src/pages/gui.tsx b/extension/react-app/src/pages/gui.tsx
index b9382bd1..4ff260fa 100644
--- a/extension/react-app/src/pages/gui.tsx
+++ b/extension/react-app/src/pages/gui.tsx
@@ -262,7 +262,6 @@ function GUI(props: GUIProps) {
 
   const onStepUserInput = (input: string, index: number) => {
     if (!client) return;
-    console.log("Sending step user input", input, index);
     client.sendStepUserInput(input, index);
   };
 
diff --git a/extension/src/activation/activate.ts b/extension/src/activation/activate.ts
index 559caf44..cd885b12 100644
--- a/extension/src/activation/activate.ts
+++ b/extension/src/activation/activate.ts
@@ -56,7 +56,7 @@ export async function activateExtension(context: vscode.ExtensionContext) {
   registerAllCodeLensProviders(context);
   registerAllCommands(context);
 
-  // Initialize IDE Protocol Client, then call "openGUI"
+  // Initialize IDE Protocol Client
   const serverUrl = getContinueServerUrl();
   ideProtocolClient = new IdeProtocolClient(
     `${serverUrl.replace("http", "ws")}/ide/ws`,
diff --git a/extension/src/continueIdeClient.ts b/extension/src/continueIdeClient.ts
index b728833f..4c1fdf1e 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,8 +14,6 @@ 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";
 
 class IdeProtocolClient {
@@ -27,17 +24,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) => {
@@ -171,7 +205,7 @@ class IdeProtocolClient {
       case "showDiff":
         this.showDiff(data.filepath, data.replacement, data.step_index);
         break;
-      case "openGUI":
+      case "getSessionId":
       case "connected":
         break;
       default:
@@ -284,10 +318,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
@@ -303,10 +333,10 @@ class IdeProtocolClient {
         }
       }, 1000);
     });
-    const resp = await this.messenger?.sendAndReceive("openGUI", {});
-    const sessionId = resp.sessionId;
+    const resp = await this.messenger?.sendAndReceive("getSessionId", {});
     // console.log("New Continue session with ID: ", sessionId);
-    return sessionId;
+    this._sessionId = resp.sessionId;
+    return resp.sessionId;
   }
 
   acceptRejectSuggestion(accept: boolean, key: SuggestionRanges) {
diff --git a/extension/src/debugPanel.ts b/extension/src/debugPanel.ts
index 487bbedf..5e1689d1 100644
--- a/extension/src/debugPanel.ts
+++ b/extension/src/debugPanel.ts
@@ -8,76 +8,6 @@ import {
 import { RangeInFile } from "./client";
 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 +57,6 @@ class WebsocketConnection {
   }
 }
 
-let streamManager = new StreamManager();
-
 export let debugPanelWebview: vscode.Webview | undefined;
 export function setupDebugPanel(
   panel: vscode.WebviewPanel | vscode.WebviewView,
@@ -147,10 +75,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 +86,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 +103,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 +145,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,15 +226,6 @@ 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();
-        break;
-      }
       case "withProgress": {
         // This message allows withProgress to be used in the webview
         if (data.done) {
diff --git a/extension/src/util/messenger.ts b/extension/src/util/messenger.ts
index b1df161b..7fd71ddd 100644
--- a/extension/src/util/messenger.ts
+++ b/extension/src/util/messenger.ts
@@ -15,6 +15,8 @@ export abstract class Messenger {
   abstract onOpen(callback: () => void): void;
 
   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);
+  }
 }