Fix: update api/websocket_manager.py

api/websocket_manager.py

@@ -1,12 +1,14 @@
 """
-WebSocket Manager — God Mode+ real-time streaming
+WebSocket Connection Manager — Production Grade
+Handles rooms, heartbeats, event buffering, reconnect support
 """
+
 import asyncio
 import json
 import time
 import uuid
+from collections import defaultdict
 from typing import Dict, List, Optional, Set
-from fastapi import WebSocket
 import structlog
 
 log = structlog.get_logger()
@@ -14,63 +16,119 @@ log = structlog.get_logger()
 
 class WebSocketManager:
     def __init__(self):
-        self._rooms: Dict[str, List[WebSocket]] = {}
-        self._chat:  Dict[str, List[WebSocket]] = {}
-
-    async def connect(self, ws: WebSocket, room: str = "global"):
-        await ws.accept()
-        if room not in self._rooms:
-            self._rooms[room] = []
-        if ws not in self._rooms[room]:
-            self._rooms[room].append(ws)
-        log.debug("WS connected", room=room)
-
-    def disconnect(self, ws: WebSocket, room: str = "global"):
-        if room in self._rooms:
-            self._rooms[room] = [w for w in self._rooms[room] if w != ws]
-        for sid, conns in self._chat.items():
-            self._chat[sid] = [w for w in conns if w != ws]
-
-    async def emit(self, task_id: str, event_type: str, data: Dict, session_id: str = ""):
-        room    = f"task:{task_id}"
-        payload = json.dumps({
-            "type":       event_type,
-            "task_id":    task_id,
-            "session_id": session_id,
-            "timestamp":  time.time(),
-            "data":       data,
-            "id":         uuid.uuid4().hex[:8],
-        })
-        await self._broadcast(room, payload)
-        await self._broadcast("logs", payload)
+        # room → set of websockets
+        self._rooms: Dict[str, Set] = defaultdict(set)
+        # ws → list of rooms
+        self._ws_rooms: Dict[object, Set[str]] = defaultdict(set)
+        # Event buffer per room (for replay on reconnect)
+        self._event_buffer: Dict[str, List] = defaultdict(list)
+        self._buffer_max = 100
+        # Active connection count
+        self._connection_count = 0
 
-    async def emit_chat(self, session_id: str, event_type: str, data: Dict):
-        room    = f"chat:{session_id}"
-        payload = json.dumps({
-            "type":       event_type,
-            "session_id": session_id,
-            "timestamp":  time.time(),
-            "data":       data,
-            "id":         uuid.uuid4().hex[:8],
+    async def connect(self, websocket, room: str):
+        await websocket.accept()
+        self._rooms[room].add(websocket)
+        self._ws_rooms[websocket].add(room)
+        self._connection_count += 1
+        log.info("WS connected", room=room, total=self._connection_count)
+
+        # Replay buffered events for this room
+        buffered = self._event_buffer.get(room, [])[-20:]
+        for event in buffered:
+            try:
+                await websocket.send_json(event)
+            except Exception:
+                pass
+
+        await websocket.send_json({
+            "type": "connected",
+            "room": room,
+            "timestamp": time.time(),
+            "buffered_events": len(buffered),
         })
-        await self._broadcast(room, payload)
 
-    async def _broadcast(self, room: str, payload: str):
-        conns = self._rooms.get(room, [])
-        dead  = []
-        for ws in conns:
+    def disconnect(self, websocket, room: Optional[str] = None):
+        if room:
+            self._rooms[room].discard(websocket)
+            self._ws_rooms[websocket].discard(room)
+        else:
+            for r in list(self._ws_rooms.get(websocket, [])):
+                self._rooms[r].discard(websocket)
+            self._ws_rooms.pop(websocket, None)
+        self._connection_count = max(0, self._connection_count - 1)
+        log.info("WS disconnected", room=room, total=self._connection_count)
+
+    async def broadcast(self, room: str, event: dict):
+        """Broadcast event to all sockets in a room."""
+        if "timestamp" not in event:
+            event["timestamp"] = time.time()
+        if "id" not in event:
+            event["id"] = str(uuid.uuid4())[:8]
+
+        # Buffer event
+        self._event_buffer[room].append(event)
+        if len(self._event_buffer[room]) > self._buffer_max:
+            self._event_buffer[room].pop(0)
+
+        dead = set()
+        for ws in list(self._rooms.get(room, [])):
             try:
-                await ws.send_text(payload)
+                await ws.send_json(event)
             except Exception:
-                dead.append(ws)
+                dead.add(ws)
+
         for ws in dead:
-            if room in self._rooms:
-                self._rooms[room] = [w for w in self._rooms[room] if w != ws]
+            self.disconnect(ws, room)
+
+    async def broadcast_global(self, event: dict):
+        """Broadcast to ALL connected websockets."""
+        for room in list(self._rooms.keys()):
+            await self.broadcast(room, event)
+
+    async def emit(self, task_id: str, event_type: str, data: dict, session_id: str = ""):
+        """Emit a structured event to a task room + logs room."""
+        event = {
+            "type": event_type,
+            "task_id": task_id,
+            "session_id": session_id,
+            "timestamp": time.time(),
+            "data": data,
+        }
+        await self.broadcast(f"task:{task_id}", event)
+        await self.broadcast("logs", event)
+        await self.broadcast("agent_status", {
+            "type": "agent_event",
+            "task_id": task_id,
+            "event_type": event_type,
+            "timestamp": time.time(),
+        })
+
+    async def emit_chat(self, session_id: str, event_type: str, data: dict):
+        """Emit event to a chat session room."""
+        event = {
+            "type": event_type,
+            "session_id": session_id,
+            "timestamp": time.time(),
+            "data": data,
+        }
+        await self.broadcast(f"chat:{session_id}", event)
 
     async def heartbeat_loop(self):
+        """Send heartbeat to all connections every 15s."""
         while True:
-            await asyncio.sleep(20)
-            hb = json.dumps({"type": "heartbeat", "timestamp": time.time()})
-            all_rooms = list(self._rooms.keys())
-            for room in all_rooms:
-                await self._broadcast(room, hb)
+            await asyncio.sleep(15)
+            heartbeat = {
+                "type": "heartbeat",
+                "timestamp": time.time(),
+                "connections": self._connection_count,
+            }
+            for room in list(self._rooms.keys()):
+                await self.broadcast(room, heartbeat)
+
+    def get_stats(self) -> dict:
+        return {
+            "total_connections": self._connection_count,
+            "rooms": {r: len(ws) for r, ws in self._rooms.items()},
+            "buffered_events": {r: len(e) for r, e in self._event_buffer.items()},
+        }