@app.websocket("/ws/text")
async def text_ws(ws: WebSocket):
    await ws.accept()
    while True:
        msg: str = await ws.receive_text()
        print(f"Text: {msg}")

@app.websocket("/ws/binary")
async def binary_ws(ws: WebSocket):
    await ws.accept()
    while True:
        data: bytes = await ws.receive_bytes()
        print(f"Received {len(data)} bytes")
        # process image, audio chunk, etc.

@app.websocket("/ws/json")
async def json_ws(ws: WebSocket):
    await ws.accept()
    while True:
        payload: dict = await ws.receive_json()
        # payload = {"action": "move", "x": 10, "y": 20}
        action = payload.get("action")
        print(f"Action: {action}")

@app.websocket("/ws/any")
async def any_ws(ws: WebSocket):
    await ws.accept()
    while True:
        message = await ws.receive()
        # message is a dict: {"type": "websocket.receive", "text": "...", "bytes": None}
        if message["type"] == "websocket.receive":
            if message.get("text"):
                print(f"Text: {message['text']}")
            elif message.get("bytes"):
                print(f"Bytes: {len(message['bytes'])} bytes")
        elif message["type"] == "websocket.disconnect":
            break

@app.websocket("/ws/demo")
async def demo_ws(ws: WebSocket):
    await ws.accept()

    # Send a plain string
    await ws.send_text("Welcome to the server!")

    # Send structured JSON (dict → serialized automatically)
    await ws.send_json({
        "type": "connected",
        "user_count": 42,
        "server_time": "2024-01-01T00:00:00Z"
    })

    # Send binary data (e.g., a small image thumbnail)
    with open("thumb.png", "rb") as f:
        await ws.send_bytes(f.read())

    # Receive and echo loop
    try:
        while True:
            msg = await ws.receive_text()
            await ws.send_json({"echo": msg, "length": len(msg)})
    except WebSocketDisconnect:
        pass

    # Graceful close with a code
    # await ws.close(code=1000)  # 1000 = normal closure

from collections import defaultdict

class RoomManager:
    def __init__(self):
        # {"general": [ws1, ws2], "dev": [ws3]}
        self.rooms: dict[str, List[WebSocket]] = defaultdict(list)

    async def join(self, room: str, websocket: WebSocket):
        await websocket.accept()
        self.rooms[room].append(websocket)
        await self.broadcast_to_room(room, f"A new user joined #{room}")

    def leave(self, room: str, websocket: WebSocket):
        if websocket in self.rooms[room]:
            self.rooms[room].remove(websocket)
        if not self.rooms[room]:
            del self.rooms[room]   # clean up empty room

    async def broadcast_to_room(self, room: str, message: str):
        for ws in self.rooms.get(room, []):
            await ws.send_text(message)

    def room_count(self, room: str) -> int:
        return len(self.rooms.get(room, []))

room_manager = RoomManager()

@app.websocket("/ws/room/{room_name}")
async def room_ws(
    websocket: WebSocket,
    room_name: str,
    username: str = Query("anonymous")
):
    await room_manager.join(room_name, websocket)
    try:
        while True:
            text = await websocket.receive_text()
            await room_manager.broadcast_to_room(
                room_name,
                f"[{room_name}] {username}: {text}"
            )
    except WebSocketDisconnect:
        room_manager.leave(room_name, websocket)
        await room_manager.broadcast_to_room(room_name, f"{username} left")

# Connect to room: ws://localhost:8000/ws/room/general?username=alice

import json
from datetime import datetime

connections: List[WebSocket] = []

async def broadcast_json(event_type: str, data: dict):
    payload = json.dumps({
        "type": event_type,
        "data": data,
        "ts": datetime.utcnow().isoformat()
    })
    dead = []
    for ws in connections:
        try:
            await ws.send_text(payload)
        except:
            dead.append(ws)  # collect failed connections
    for ws in dead:
        connections.remove(ws)  # clean up dead ones

# Trigger a broadcast from an HTTP endpoint (e.g., admin pushes an alert)
@app.post("/admin/broadcast")
async def admin_broadcast(message: str):
    await broadcast_json("admin_alert", {"message": message})
    return {"sent_to": len(connections)}

# pip install redis asyncio-redis or redis[asyncio]
import asyncio
import json
import redis.asyncio as aioredis

redis_client = None
local_connections: List[WebSocket] = []

@asynccontextmanager
async def lifespan(app):
    global redis_client
    redis_client = aioredis.from_url("redis://localhost:6379")
    # Start subscriber background task
    asyncio.create_task(redis_subscriber())
    yield
    await redis_client.close()

async def redis_subscriber():
    "Subscribe to Redis, relay to local WebSocket clients"
    pubsub = redis_client.pubsub()
    await pubsub.subscribe("chat")
    async for message in pubsub.listen():
        if message["type"] == "message":
            text = message["data"].decode()
            for ws in local_connections:
                await ws.send_text(text)

@app.websocket("/ws")
async def ws_endpoint(ws: WebSocket):
    await ws.accept()
    local_connections.append(ws)
    try:
        while True:
            msg = await ws.receive_text()
            # Publish to Redis → ALL workers relay to their clients
            await redis_client.publish("chat", msg)
    except WebSocketDisconnect:
        local_connections.remove(ws)

@app.websocket("/ws/secure")
async def secure_ws(websocket: WebSocket):
    await websocket.accept()   # accept first (connection is open)

    # Wait for the first message — must be an auth token
    auth_msg = await websocket.receive_json()
    token = auth_msg.get("token")

    user = get_user_from_token(token)
    if not user:
        await websocket.send_json({"type": "error", "msg": "Unauthorized"})
        await websocket.close(code=1008)
        return

    await websocket.send_json({"type": "auth_ok"})

    try:
        while True:
            data = await websocket.receive_json()
            await websocket.send_json({"echo": data})
    except WebSocketDisconnect:
        pass

// JS client:
// ws.onopen = () => ws.send(JSON.stringify({token: "eyJhbGci..."}));

from fastapi import Cookie

@app.websocket("/ws")
async def cookie_ws(
    websocket: WebSocket,
    session: str = Cookie(None)   # reads the "session" cookie
):
    if not session:
        await websocket.close(code=1008)
        return
    user = validate_session(session)  # your session lookup
    await websocket.accept()
    await websocket.send_json({"type": "auth_ok", "user": user})