release: bump version to 0.48.0

backend/app/core/websocket/manager.py

@@ -1,9 +1,6 @@
 """WebSocket Connection Manager"""
 
-import json
-import asyncio
 from typing import Dict, Set, Optional
-from datetime import datetime
 from fastapi import WebSocket
 import redis.asyncio as redis
 
@@ -15,6 +12,8 @@ class ConnectionManager:
 
     def __init__(self):
         self.active_connections: Dict[str, Set[WebSocket]] = {}  # user_id -> connections
+        self.channel_subscriptions: Dict[str, Set[WebSocket]] = {}
+        self.websocket_channels: Dict[WebSocket, Set[str]] = {}
         self.redis_client: Optional[redis.Redis] = None
 
     async def connect(self, websocket: WebSocket, user_id: str):
@@ -40,6 +39,39 @@ class ConnectionManager:
             self.active_connections[user_id].discard(websocket)
             if not self.active_connections[user_id]:
                 del self.active_connections[user_id]
+        self.unsubscribe_all(websocket)
+
+    def subscribe(self, websocket: WebSocket, channels: list[str]):
+        normalized_channels = {
+            str(channel).strip()
+            for channel in channels
+            if str(channel).strip()
+        }
+        if not normalized_channels:
+            return
+
+        socket_channels = self.websocket_channels.setdefault(websocket, set())
+        for channel in normalized_channels:
+            self.channel_subscriptions.setdefault(channel, set()).add(websocket)
+            socket_channels.add(channel)
+
+    def unsubscribe(self, websocket: WebSocket, channels: list[str]):
+        for channel in {str(channel).strip() for channel in channels if str(channel).strip()}:
+            subscribers = self.channel_subscriptions.get(channel)
+            if subscribers is not None:
+                subscribers.discard(websocket)
+                if not subscribers:
+                    del self.channel_subscriptions[channel]
+            socket_channels = self.websocket_channels.get(websocket)
+            if socket_channels is not None:
+                socket_channels.discard(channel)
+                if not socket_channels:
+                    del self.websocket_channels[websocket]
+
+    def unsubscribe_all(self, websocket: WebSocket):
+        channels = list(self.websocket_channels.get(websocket, set()))
+        if channels:
+            self.unsubscribe(websocket, channels)
 
     async def send_personal_message(self, message: dict, user_id: str):
         if user_id in self.active_connections:
@@ -54,13 +86,19 @@ class ConnectionManager:
             for user_id in self.active_connections:
                 await self.send_personal_message(message, user_id)
         else:
-            await self.send_personal_message(message, channel)
+            for connection in list(self.channel_subscriptions.get(channel, set())):
+                try:
+                    await connection.send_json(message)
+                except Exception:
+                    self.unsubscribe_all(connection)
 
     async def close_all(self):
         for user_id in self.active_connections:
             for connection in self.active_connections[user_id]:
                 await connection.close()
         self.active_connections.clear()
+        self.channel_subscriptions.clear()
+        self.websocket_channels.clear()
 
 
 manager = ConnectionManager()