10 files changed, 70 insertions, 14 deletions
diff --git a/continuedev/src/continuedev/server/gui.py b/continuedev/src/continuedev/server/gui.py
index c0957395..9bb621cb 100644
--- a/continuedev/src/continuedev/server/gui.py
+++ b/continuedev/src/continuedev/server/gui.py
@@ -32,10 +32,6 @@ class AppStatus:
 Server.handle_exit = AppStatus.handle_exit
 
 
-async def session(x_continue_session_id: str = Header("anonymous")) -> Session:
-    return await session_manager.get_session(x_continue_session_id)
-
-
 async def websocket_session(session_id: str) -> Session:
     return await session_manager.get_session(session_id)
 
diff --git a/continuedev/src/continuedev/server/ide.py b/continuedev/src/continuedev/server/ide.py
index 87374928..e4c2d714 100644
--- a/continuedev/src/continuedev/server/ide.py
+++ b/continuedev/src/continuedev/server/ide.py
@@ -482,6 +482,7 @@ async def websocket_endpoint(websocket: WebSocket, session_id: str = None):
             "error_title": e.__str__() or e.__repr__(), "error_message": '\n'.join(traceback.format_exception(e))})
         raise e
     finally:
+        print("Closing ide websocket")
         if websocket.client_state != WebSocketState.DISCONNECTED:
             await websocket.close()
 
diff --git a/continuedev/src/continuedev/server/main.py b/continuedev/src/continuedev/server/main.py
index b92c9fa3..d78eaecf 100644
--- a/continuedev/src/continuedev/server/main.py
+++ b/continuedev/src/continuedev/server/main.py
@@ -59,12 +59,17 @@ LOG_CONFIG = {
         'handlers': ['file']
     }
 }
+
 print(f"Log path: {log_path}")
+sys.stdout = open(log_path, "a")
+sys.stderr = open(log_path, "a")
+print("Testing logs")
 
 
 def run_server():
-    config = uvicorn.Config(app, host="0.0.0.0",
-                            port=args.port, log_config=LOG_CONFIG)
+    config = uvicorn.Config(app, host="0.0.0.0", port=args.port,
+                            # log_config=LOG_CONFIG
+                            )
     server = uvicorn.Server(config)
 
     server.run()
diff --git a/continuedev/src/continuedev/server/session_manager.py b/continuedev/src/continuedev/server/session_manager.py
index 3136f1bf..4e47098a 100644
--- a/continuedev/src/continuedev/server/session_manager.py
+++ b/continuedev/src/continuedev/server/session_manager.py
@@ -4,6 +4,8 @@ from typing import Any, Dict, List, Union
 from uuid import uuid4
 import json
 
+from fastapi.websockets import WebSocketState
+
 from ..libs.util.paths import getSessionFilePath, getSessionsFolderPath
 from ..models.filesystem_edit import FileEditWithFullContents
 from ..libs.constants.main import CONTINUE_SESSIONS_FOLDER
@@ -59,6 +61,8 @@ class SessionManager:
         return self.sessions[session_id]
 
     async def new_session(self, ide: AbstractIdeProtocolServer, session_id: Union[str, None] = None) -> Session:
+        print("New session: ", session_id)
+
         full_state = None
         if session_id is not None and os.path.exists(getSessionFilePath(session_id)):
             with open(getSessionFilePath(session_id), "r") as f:
@@ -82,13 +86,16 @@ class SessionManager:
         return session
 
     def remove_session(self, session_id: str):
-        del self.sessions[session_id]
+        print("Removing session: ", session_id)
+        if session_id in self.sessions:
+            ws_to_close = self.sessions[session_id].ide.websocket
+            if ws_to_close is not None and ws_to_close.client_state != WebSocketState.DISCONNECTED:
+                self.sessions[session_id].autopilot.ide.websocket.close()
+            del self.sessions[session_id]
 
     async def persist_session(self, session_id: str):
         """Save the session's FullState as a json file"""
         full_state = await 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)
 
diff --git a/extension/package-lock.json b/extension/package-lock.json
index fcd97d50..b210256b 100644
--- a/extension/package-lock.json
+++ b/extension/package-lock.json
@@ -1,12 +1,12 @@
 {
   "name": "continue",
-  "version": "0.0.213",
+  "version": "0.0.214",
   "lockfileVersion": 2,
   "requires": true,
   "packages": {
     "": {
       "name": "continue",
-      "version": "0.0.213",
+      "version": "0.0.214",
       "license": "Apache-2.0",
       "dependencies": {
         "@electron/rebuild": "^3.2.10",
diff --git a/extension/package.json b/extension/package.json
index e1004a8c..72cfa6ab 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.213",
+  "version": "0.0.214",
   "publisher": "Continue",
   "engines": {
     "vscode": "^1.67.0"
diff --git a/extension/react-app/src/hooks/ContinueGUIClientProtocol.ts b/extension/react-app/src/hooks/ContinueGUIClientProtocol.ts
index b8019664..5a5d4c30 100644
--- a/extension/react-app/src/hooks/ContinueGUIClientProtocol.ts
+++ b/extension/react-app/src/hooks/ContinueGUIClientProtocol.ts
@@ -17,6 +17,13 @@ class ContinueGUIClientProtocol extends AbstractContinueGUIClientProtocol {
     this.messenger = useVscodeMessagePassing
       ? new VscodeMessenger(serverUrlWithSessionId)
       : new WebsocketMessenger(serverUrlWithSessionId);
+
+    this.messenger.onClose(() => {
+      console.log("GUI -> IDE websocket closed");
+    });
+    this.messenger.onError((error) => {
+      console.log("GUI -> IDE websocket error", error);
+    });
   }
 
   sendMainInput(input: string) {
diff --git a/extension/react-app/src/hooks/messenger.ts b/extension/react-app/src/hooks/messenger.ts
index 00ce1fbb..ecf646c7 100644
--- a/extension/react-app/src/hooks/messenger.ts
+++ b/extension/react-app/src/hooks/messenger.ts
@@ -13,6 +13,8 @@ export abstract class Messenger {
   abstract onClose(callback: () => void): void;
 
   abstract sendAndReceive(messageType: string, data: any): Promise<any>;
+
+  abstract onError(callback: (error: any) => void): void;
 }
 
 export class WebsocketMessenger extends Messenger {
@@ -20,6 +22,7 @@ export class WebsocketMessenger extends Messenger {
   private onMessageListeners: {
     [messageType: string]: ((data: object) => void)[];
   } = {};
+  private onErrorListeners: ((error: any) => void)[] = [];
   private onOpenListeners: (() => void)[] = [];
   private onCloseListeners: (() => void)[] = [];
   private serverUrl: string;
@@ -37,6 +40,9 @@ export class WebsocketMessenger extends Messenger {
         this.onMessageType(messageType, listener);
       }
     }
+    for (const listener of this.onErrorListeners) {
+      this.onError(listener);
+    }
     return newWebsocket;
   }
 
@@ -95,4 +101,8 @@ export class WebsocketMessenger extends Messenger {
   onClose(callback: () => void): void {
     this.websocket.addEventListener("close", callback);
   }
+
+  onError(callback: (error: any) => void): void {
+    this.websocket.addEventListener("error", callback);
+  }
 }
diff --git a/extension/react-app/src/hooks/vscodeMessenger.ts b/extension/react-app/src/hooks/vscodeMessenger.ts
index ba19586b..13f5092b 100644
--- a/extension/react-app/src/hooks/vscodeMessenger.ts
+++ b/extension/react-app/src/hooks/vscodeMessenger.ts
@@ -38,6 +38,14 @@ export class VscodeMessenger extends Messenger {
     });
   }
 
+  onError(callback: (error: any) => void): void {
+    window.addEventListener("message", (event: any) => {
+      if (event.data.type === "websocketForwardingError") {
+        callback(event.data.error);
+      }
+    });
+  }
+
   sendAndReceive(messageType: string, data: any): Promise<any> {
     return new Promise((resolve) => {
       const handler = (event: any) => {
diff --git a/extension/src/debugPanel.ts b/extension/src/debugPanel.ts
index 3c4f8481..4785ba20 100644
--- a/extension/src/debugPanel.ts
+++ b/extension/src/debugPanel.ts
@@ -17,17 +17,20 @@ class WebsocketConnection {
   private _onMessage: (message: string) => void;
   private _onOpen: () => void;
   private _onClose: () => void;
+  private _onError: (e: any) => void;
 
   constructor(
     url: string,
     onMessage: (message: string) => void,
     onOpen: () => void,
-    onClose: () => void
+    onClose: () => void,
+    onError: (e: any) => void
   ) {
     this._ws = new WebSocket(url);
     this._onMessage = onMessage;
     this._onOpen = onOpen;
     this._onClose = onClose;
+    this._onError = onError;
 
     this._ws.addEventListener("message", (event) => {
       this._onMessage(event.data);
@@ -38,6 +41,9 @@ class WebsocketConnection {
     this._ws.addEventListener("open", () => {
       this._onOpen();
     });
+    this._ws.addEventListener("error", (e: any) => {
+      this._onError(e);
+    });
   }
 
   public send(message: string) {
@@ -147,12 +153,20 @@ export function setupDebugPanel(
           url,
         });
       };
+      const onError = (e: any) => {
+        panel.webview.postMessage({
+          type: "websocketForwardingError",
+          url,
+          error: e,
+        });
+      };
       try {
         const connection = new WebsocketConnection(
           url,
           onMessage,
           onOpen,
-          onClose
+          onClose,
+          onError
         );
         websocketConnections[url] = connection;
         resolve(null);
@@ -198,6 +212,14 @@ export function setupDebugPanel(
         if (typeof websocketConnections[url] === "undefined") {
           await connectWebsocket(url);
         }
+        console.log(
+          "Websocket connection requested by GUI already open at",
+          url
+        );
+        panel.webview.postMessage({
+          type: "websocketForwardingOpen",
+          url,
+        });
         break;
       }
       case "websocketForwardingMessage": {