refactor: :goal_net: log errors from websocket, better logging from uvicorn

4 files changed, 18 insertions, 9 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)