fix: remove notification manager singleton

Thymis-io · Dec 6, 2024 · af3782b · af3782b
1 parent 0bc9bd1
commit af3782b
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 7 deletions.
diff --git 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,14 +137,15 @@ 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()
     logger.info("frontend started")
     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
@@ -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
@@ -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"])