stuff

2 files changed, 4 insertions, 2 deletions

diff --git a/continuedev/src/continuedev/server/gui.py b/continuedev/src/continuedev/server/gui.py
index 70778dfe..fa203c28 100644
--- a/continuedev/src/continuedev/server/gui.py
+++ b/continuedev/src/continuedev/server/gui.py
@@ -179,7 +179,7 @@ async def websocket_endpoint(websocket: WebSocket, session: Session = Depends(we
 
         while AppStatus.should_exit is False:
             message = await websocket.receive_text()
-            print("Received message", message)
+            print("Received GUI message", message)
             if type(message) is str:
                 message = json.loads(message)
 
diff --git a/continuedev/src/continuedev/server/ide.py b/continuedev/src/continuedev/server/ide.py
index aa69a6cb..d6a28c92 100644
--- a/continuedev/src/continuedev/server/ide.py
+++ b/continuedev/src/continuedev/server/ide.py
@@ -444,6 +444,7 @@ async def websocket_endpoint(websocket: WebSocket, session_id: str = None):
             message_type = message["messageType"]
             data = message["data"]
 
+            print("Received IDE message: ", message_type)
             create_async_task(
                 ideProtocolServer.handle_json(message_type, data))
 
@@ -475,4 +476,5 @@ async def websocket_endpoint(websocket: WebSocket, session_id: str = None):
 
         posthog_logger.capture_event("session_ended", {
             "session_id": ideProtocolServer.session_id})
-        session_manager.registered_ides.pop(ideProtocolServer.session_id)
+        if ideProtocolServer.session_id in session_manager.registered_ides:
+            session_manager.registered_ides.pop(ideProtocolServer.session_id)