diff options
Diffstat (limited to 'continuedev/src')
| -rw-r--r-- | continuedev/src/continuedev/server/gui.py | 2 | ||||
| -rw-r--r-- | continuedev/src/continuedev/server/session_manager.py | 10 |
2 files changed, 7 insertions, 5 deletions
diff --git a/continuedev/src/continuedev/server/gui.py b/continuedev/src/continuedev/server/gui.py index 9bb621cb..58e875b9 100644 --- a/continuedev/src/continuedev/server/gui.py +++ b/continuedev/src/continuedev/server/gui.py @@ -198,4 +198,4 @@ async def websocket_endpoint(websocket: WebSocket, session: Session = Depends(we await websocket.close() await session_manager.persist_session(session.session_id) - session_manager.remove_session(session.session_id) + await session_manager.remove_session(session.session_id) diff --git a/continuedev/src/continuedev/server/session_manager.py b/continuedev/src/continuedev/server/session_manager.py index 4e47098a..511dc070 100644 --- a/continuedev/src/continuedev/server/session_manager.py +++ b/continuedev/src/continuedev/server/session_manager.py @@ -85,12 +85,14 @@ class SessionManager: create_async_task(autopilot.run_policy()) return session - def remove_session(self, session_id: str): + async def remove_session(self, session_id: str): 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() + if session_id in self.registered_ides: + ws_to_close = self.registered_ides[session_id].websocket + if ws_to_close is not None and ws_to_close.client_state != WebSocketState.DISCONNECTED: + await self.sessions[session_id].autopilot.ide.websocket.close() + del self.sessions[session_id] async def persist_session(self, session_id: str): |