Update main.py

main.py

@@ -2,9 +2,10 @@ from fastapi import FastAPI, WebSocket, WebSocketDisconnect
 from fastapi.responses import HTMLResponse
 from pathlib import Path
 import uvicorn
+import json
 
 app = FastAPI()
-rooms: dict[str, set[WebSocket]] = {}
+rooms = {}
 
 @app.get("/")
 async def index():
@@ -12,32 +13,87 @@ async def index():
 
 @app.get("/ping")
 async def ping():
-    return {"status": "alive", "rooms": {k: len(v) for k, v in rooms.items()}}
+    return {"status": "alive"}
 
 @app.websocket("/ws/{room}")
-async def ws_relay(websocket: WebSocket, room: str):
+async def ws_endpoint(websocket: WebSocket, room: str):
     await websocket.accept()
-    rooms.setdefault(room, set()).add(websocket)
-    print(f"[+] Join room '{room}' | total: {len(rooms[room])}")
+    if room not in rooms:
+        rooms[room] = {"peers": set(), "server": None}
+    r = rooms[room]
+    r["peers"].add(websocket)
+    await broadcast_state(r)
     try:
         while True:
-            data = await websocket.receive_bytes()
-            if len(data) < 100:
-                continue  # skip keep-alive ping
-            dead = []
-            for peer in rooms.get(room, set()):
-                if peer != websocket:
-                    try:
-                        await peer.send_bytes(data)
-                    except Exception:
-                        dead.append(peer)
-            for d in dead:
-                rooms[room].discard(d)
+            msg = await websocket.receive()
+            if "text" in msg:
+                await handle_signal(r, websocket, json.loads(msg["text"]))
+            elif "bytes" in msg:
+                raw = msg["bytes"]
+                if len(raw) < 2:
+                    continue
+                tag = raw[0]
+                if tag == 1:
+                    await relay(r, websocket, raw)
+                elif tag == 2 and websocket == r["server"]:
+                    await relay(r, websocket, raw)
     except (WebSocketDisconnect, Exception):
-        rooms.get(room, set()).discard(websocket)
-        if room in rooms and not rooms[room]:
-            del rooms[room]
-        print(f"[-] Left room '{room}'")
+        r["peers"].discard(websocket)
+        if websocket == r["server"]:
+            r["server"] = None
+        if not r["peers"]:
+            rooms.pop(room, None)
+        else:
+            await broadcast_state(r)
+
+async def relay(r, sender, data):
+    dead = []
+    for peer in list(r["peers"]):
+        if peer != sender:
+            try:
+                await peer.send_bytes(data)
+            except:
+                dead.append(peer)
+    for d in dead:
+        r["peers"].discard(d)
+
+async def handle_signal(r, ws, data):
+    t = data.get("type")
+    if t == "set_server":
+        old = r["server"]
+        r["server"] = ws
+        if old and old != ws and old in r["peers"]:
+            try:
+                await old.send_text(json.dumps({"type": "force_client"}))
+            except:
+                pass
+        await broadcast_state(r)
+    elif t == "set_client":
+        if ws == r["server"]:
+            r["server"] = None
+        await broadcast_state(r)
+    elif t == "remote_speaker":
+        srv = r["server"]
+        if srv and srv in r["peers"]:
+            try:
+                await srv.send_text(json.dumps({
+                    "type": "remote_speaker",
+                    "enabled": data.get("enabled", True)
+                }))
+            except:
+                pass
+
+async def broadcast_state(r):
+    for peer in list(r["peers"]):
+        try:
+            await peer.send_text(json.dumps({
+                "type": "state",
+                "peers": len(r["peers"]),
+                "has_server": r["server"] is not None,
+                "is_server": peer == r["server"]
+            }))
+        except:
+            r["peers"].discard(peer)
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)