manager.py 2.4 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667
  1. import asyncio
  2. import time
  3. from typing import Optional
  4. from fastapi import WebSocket
  5. from opendevin.core.config import AppConfig
  6. from opendevin.core.logger import opendevin_logger as logger
  7. from .session import Session
  8. class SessionManager:
  9. _sessions: dict[str, Session] = {}
  10. cleanup_interval: int = 300
  11. session_timeout: int = 600
  12. def __init__(self, config: AppConfig):
  13. asyncio.create_task(self._cleanup_sessions())
  14. self.config = config
  15. def add_or_restart_session(self, sid: str, ws_conn: WebSocket) -> Session:
  16. if sid in self._sessions:
  17. asyncio.create_task(self._sessions[sid].close())
  18. self._sessions[sid] = Session(sid=sid, ws=ws_conn, config=self.config)
  19. return self._sessions[sid]
  20. def get_session(self, sid: str) -> Session | None:
  21. if sid not in self._sessions:
  22. return None
  23. return self._sessions.get(sid)
  24. async def send(self, sid: str, data: dict[str, object]) -> bool:
  25. """Sends data to the client."""
  26. if sid not in self._sessions:
  27. return False
  28. return await self._sessions[sid].send(data)
  29. async def send_error(self, sid: str, message: str) -> bool:
  30. """Sends an error message to the client."""
  31. return await self.send(sid, {'error': True, 'message': message})
  32. async def send_message(self, sid: str, message: str) -> bool:
  33. """Sends a message to the client."""
  34. return await self.send(sid, {'message': message})
  35. async def _cleanup_sessions(self):
  36. while True:
  37. current_time = time.time()
  38. session_ids_to_remove = []
  39. for sid, session in list(self._sessions.items()):
  40. # if session inactive for a long time, remove it
  41. if (
  42. not session.is_alive
  43. and current_time - session.last_active_ts > self.session_timeout
  44. ):
  45. session_ids_to_remove.append(sid)
  46. for sid in session_ids_to_remove:
  47. to_del_session: Optional[Session] = self._sessions.pop(sid, None)
  48. if to_del_session is not None:
  49. await to_del_session.close()
  50. logger.info(
  51. f'Session {sid} and related resource have been removed due to inactivity.'
  52. )
  53. await asyncio.sleep(self.cleanup_interval)