ws_router.py

from fastapi import APIRouter, WebSocket, WebSocketDisconnect
import asyncio
import json
from datetime import datetime

router = APIRouter()
clients = set()

@router.websocket("/ws")
async def codex_stream(websocket: WebSocket):
    await websocket.accept()
    clients.add(websocket)
    await websocket.send_text(json.dumps({
        "type": "system",
        "message": "CodexMesh stream connected",
        "time": datetime.utcnow().isoformat()
    }))

    try:
        while True:
            await asyncio.sleep(1)
    except WebSocketDisconnect:
        clients.remove(websocket)

async def broadcast(event: dict):
    dead = set()
    for ws in clients:
        try:
            await ws.send_text(json.dumps(event))
        except Exception:
            dead.add(ws)

    for ws in dead:
        clients.remove(ws)