diff --git a/controller/thymis_controller/main.py b/controller/thymis_controller/main.py index c1ce93ce..b0455403 100644 --- a/controller/thymis_controller/main.py +++ b/controller/thymis_controller/main.py @@ -16,7 +16,7 @@ from fastapi.middleware.cors import CORSMiddleware from fastapi.middleware.gzip import GZipMiddleware from thymis_controller.config import global_settings -from thymis_controller.notifications import notification_manager +from thymis_controller.notifications import NotificationManager from thymis_controller.routers import agent, api, auth, frontend logger = logging.getLogger(__name__) @@ -137,6 +137,7 @@ async def lifespan(app: FastAPI): peform_db_upgrade() init_password_file() init_ssh_key() + notification_manager = NotificationManager() notification_manager.start() logger.info("starting frontend") await frontend.frontend.run() @@ -144,7 +145,7 @@ async def lifespan(app: FastAPI): asyncio.get_event_loop().create_task(frontend.frontend.raise_if_terminated()) logger.info("frontend raise_if_terminated task created") logger.info("Starting controller at \033[1m%s\033[0m", global_settings.BASE_URL) - yield + yield {"notification_manager": notification_manager} notification_manager.stop() logger.info("stopping frontend") await frontend.frontend.stop() diff --git a/controller/thymis_controller/notifications.py b/controller/thymis_controller/notifications.py index 54cda268..d87519d5 100644 --- a/controller/thymis_controller/notifications.py +++ b/controller/thymis_controller/notifications.py @@ -110,6 +110,3 @@ async def _broadcast(self, message: Notification): except Exception: self.disconnect(connection) traceback.print_exc() - - -notification_manager = NotificationManager() diff --git a/controller/thymis_controller/routers/api.py b/controller/thymis_controller/routers/api.py index f72cf404..c5dbbae8 100644 --- a/controller/thymis_controller/routers/api.py +++ b/controller/thymis_controller/routers/api.py @@ -14,7 +14,6 @@ require_valid_user_session, ) from thymis_controller.models.state import State -from thymis_controller.notifications import notification_manager from thymis_controller.routers import task from thymis_controller.tcp_to_ws import ( channel_to_websocket, @@ -148,7 +147,7 @@ def download_image( @router.websocket("/notification") async def notification_websocket(websocket: WebSocket): - await notification_manager.connect(websocket) + await websocket.state.notification_manager.connect(websocket) @router.get("/history", tags=["history"])