3 files changed, 59 insertions, 46 deletions
diff --git a/continuedev/src/continuedev/server/gui.py b/continuedev/src/continuedev/server/gui.py
index cc6235e9..c0178920 100644
--- a/continuedev/src/continuedev/server/gui.py
+++ b/continuedev/src/continuedev/server/gui.py
@@ -130,29 +130,36 @@ class GUIProtocolServer(AbstractGUIProtocolServer):
 
 @router.websocket("/ws")
 async def websocket_endpoint(websocket: WebSocket, session: Session = Depends(websocket_session)):
-    await websocket.accept()
-
-    print("Session started")
-    session_manager.register_websocket(session.session_id, websocket)
-    protocol = GUIProtocolServer(session)
-    protocol.websocket = websocket
-
-    # Update any history that may have happened before connection
-    await protocol.send_available_slash_commands()
-    await protocol.send_state_update()
-
-    while AppStatus.should_exit is False:
-        message = await websocket.receive_text()
-        print("Received message", message)
-        if type(message) is str:
-            message = json.loads(message)
-
-        if "messageType" not in message or "data" not in message:
-            continue
-        message_type = message["messageType"]
-        data = message["data"]
-
-        protocol.handle_json(message_type, data)
-
-    print("Closing websocket")
-    await websocket.close()
+    try:
+        print("Received websocket connection at url: ", websocket.url)
+        await websocket.accept()
+
+        print("Session started")
+        session_manager.register_websocket(session.session_id, websocket)
+        protocol = GUIProtocolServer(session)
+        protocol.websocket = websocket
+
+        # Update any history that may have happened before connection
+        await protocol.send_available_slash_commands()
+        # await protocol.send_state_update() THIS WAS CAUSING A LOT OF ISSUES. Don't uncomment.
+
+        while AppStatus.should_exit is False:
+            message = await websocket.receive_text()
+            print("Received message", message)
+            if type(message) is str:
+                message = json.loads(message)
+
+            if "messageType" not in message or "data" not in message:
+                continue
+            message_type = message["messageType"]
+            data = message["data"]
+
+            protocol.handle_json(message_type, data)
+
+    except Exception as e:
+        print("ERROR in gui websocket: ", e)
+        raise e
+    finally:
+        print("Closing gui websocket")
+        await websocket.close()
+        session_manager.remove_session(session.session_id)
diff --git a/continuedev/src/continuedev/server/ide.py b/continuedev/src/continuedev/server/ide.py
index 65f3ee74..3c3555f1 100644
--- a/continuedev/src/continuedev/server/ide.py
+++ b/continuedev/src/continuedev/server/ide.py
@@ -349,22 +349,28 @@ class IdeProtocolServer(AbstractIdeProtocolServer):
 
 @router.websocket("/ws")
 async def websocket_endpoint(websocket: WebSocket):
-    await websocket.accept()
-    print("Accepted websocket connection from, ", websocket.client)
-    await websocket.send_json({"messageType": "connected", "data": {}})
-
-    ideProtocolServer = IdeProtocolServer(session_manager)
-    ideProtocolServer.websocket = websocket
-
-    while True:
-        message = await websocket.receive_text()
-        message = json.loads(message)
-
-        if "messageType" not in message or "data" not in message:
-            continue
-        message_type = message["messageType"]
-        data = message["data"]
-
-        await ideProtocolServer.handle_json(message_type, data)
-
-    await websocket.close()
+    try:
+        await websocket.accept()
+        print("Accepted websocket connection from, ", websocket.client)
+        await websocket.send_json({"messageType": "connected", "data": {}})
+
+        ideProtocolServer = IdeProtocolServer(session_manager)
+        ideProtocolServer.websocket = websocket
+
+        while AppStatus.should_exit is False:
+            message = await websocket.receive_text()
+            message = json.loads(message)
+
+            if "messageType" not in message or "data" not in message:
+                continue
+            message_type = message["messageType"]
+            data = message["data"]
+
+            await ideProtocolServer.handle_json(message_type, data)
+
+        print("Closing ide websocket")
+        await websocket.close()
+    except Exception as e:
+        print("Error in ide websocket: ", e)
+        await websocket.close()
+        raise e
diff --git a/continuedev/src/continuedev/server/session_manager.py b/continuedev/src/continuedev/server/session_manager.py
index ebea08a5..99a38146 100644
--- a/continuedev/src/continuedev/server/session_manager.py
+++ b/continuedev/src/continuedev/server/session_manager.py
@@ -42,7 +42,7 @@ class SessionManager:
 
     def get_session(self, session_id: str) -> Session:
         if session_id not in self.sessions:
-            raise KeyError("Session ID not recognized")
+            raise KeyError("Session ID not recognized", session_id)
         return self.sessions[session_id]
 
     def new_session(self, ide: AbstractIdeProtocolServer) -> str: