God Mode+ v3 fix: api/websocket_manager.py

api/websocket_manager.py

@@ -1,14 +1,12 @@
 """
-WebSocket Connection Manager — Production Grade
-Handles rooms, heartbeats, event buffering, reconnect support
+WebSocket Manager — God Mode+ real-time streaming
 """
-
 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()
@@ -16,119 +14,63 @@ log = structlog.get_logger()
 
 class WebSocketManager:
     def __init__(self):
-        # 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 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),
+        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)
 
-    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)
+    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],
+        })
+        await self._broadcast(room, payload)
 
-        dead = set()
-        for ws in list(self._rooms.get(room, [])):
+    async def _broadcast(self, room: str, payload: str):
+        conns = self._rooms.get(room, [])
+        dead  = []
+        for ws in conns:
             try:
-                await ws.send_json(event)
+                await ws.send_text(payload)
             except Exception:
-                dead.add(ws)
-
+                dead.append(ws)
         for ws in dead:
-            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)
+            if room in self._rooms:
+                self._rooms[room] = [w for w in self._rooms[room] if w != ws]
 
     async def heartbeat_loop(self):
-        """Send heartbeat to all connections every 15s."""
         while True:
-            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()},
-        }
+            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)