FastAPI WebSockets — Real-time Communication

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from typing import Dict
import asyncio

app = FastAPI()


class ForgeSocketManager:
    """Manages active WebSocket connections keyed by user ID.

    The dictionary is the single source of truth for who is currently
    connected. Every code path that touches it must respect the
    connect/disconnect contract — no exceptions.
    """

    def __init__(self):
        # Keyed by user_id so we can route messages without iterating everything
        self.active_connections: Dict[str, WebSocket] = {}

    async def connect(self, user_id: str, websocket: WebSocket) -> None:
        await websocket.accept()
        self.active_connections[user_id] = websocket

    def disconnect(self, user_id: str) -> None:
        # .pop() with a default is safer than del — avoids KeyError on double-disconnect
        self.active_connections.pop(user_id, None)

    async def send_personal_message(self, message: str, user_id: str) -> None:
        ws = self.active_connections.get(user_id)
        if ws is not None:
            await ws.send_text(message)

    async def broadcast(self, message: str) -> None:
        # asyncio.gather fans out all sends concurrently instead of awaiting each one
        # return_exceptions=True prevents one failed send from cancelling the rest
        await asyncio.gather(
            *[ws.send_text(message) for ws in self.active_connections.values()],
            return_exceptions=True,
        )


manager = ForgeSocketManager()


@app.websocket("/ws/{user_id}")
async def websocket_endpoint(websocket: WebSocket, user_id: str):
    await manager.connect(user_id, websocket)
    try:
        while True:
            data = await websocket.receive_text()
            await manager.broadcast(f"User {user_id}: {data}")
    except WebSocketDisconnect:
        # Clean close frame received — remove from manager and notify others
        manager.disconnect(user_id)
        await manager.broadcast(f"User {user_id} has left the forge.")

from fastapi import FastAPI, WebSocket, Query, status
from jose import JWTError, jwt
from datetime import datetime

app = FastAPI()

SECRET_KEY = "your-secret-key"  # Load from environment in production
ALGORITHM = "HS256"


def decode_ws_token(token: str) -> dict | None:
    """Returns the decoded payload or None if the token is invalid or expired."""
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        # Verify expiry explicitly — some JWT libraries are lenient about this
        if payload.get("exp") and datetime.utcnow().timestamp() > payload["exp"]:
            return None
        return payload
    except JWTError:
        return None


@app.websocket("/secure-ws")
async def secure_socket(websocket: WebSocket, token: str = Query(...)):
    # Validate BEFORE accept() — this is the hard rule
    payload = decode_ws_token(token)
    if payload is None:
        # 1008 = Policy Violation — the correct code for auth failure
        # 1000 = Normal Closure — wrong signal; implies success
        await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
        return

    user_id: str = payload.get("sub")
    await websocket.accept()
    await websocket.send_text(f"Authenticated. Welcome, {user_id}.")

    try:
        while True:
            data = await websocket.receive_text()
            await websocket.send_text(f"Echo: {data}")
    except WebSocketDisconnect:
        pass  # Clean up handled by ConnectionManager in the full implementation

import asyncio
import redis.asyncio as aioredis
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from contextlib import asynccontextmanager


local_connections: dict[str, WebSocket] = {}
CHANNEL = "forge:broadcast"
redis: aioredis.Redis | None = None


async def redis_listener(r: aioredis.Redis) -> None:
    """Long-running coroutine: subscribes to Redis and forwards to local clients."""
    pubsub = r.pubsub()
    await pubsub.subscribe(CHANNEL)
    async for message in pubsub.listen():
        if message["type"] != "message":
            continue
        data: str = message["data"].decode()
        # Collect stale connections during the fan-out rather than modifying
        # the dict mid-iteration, which raises RuntimeError in Python 3.x
        stale: list[str] = []
        await asyncio.gather(
            *[
                ws.send_text(data)
                for uid, ws in local_connections.items()
            ],
            return_exceptions=True,
        )
        # Clean up any connections that raised during the gather
        for uid in stale:
            local_connections.pop(uid, None)


@asynccontextmanager
async def lifespan(app: FastAPI):
    global redis
    redis = aioredis.from_url("redis://localhost:6379", decode_responses=False)
    # Start the listener as a background task tied to the worker's lifespan
    listener_task = asyncio.create_task(redis_listener(redis))
    yield
    # Graceful shutdown — cancel the listener and close the Redis connection
    listener_task.cancel()
    await redis.aclose()


app = FastAPI(lifespan=lifespan)


@app.websocket("/ws/{user_id}")
async def ws_endpoint(websocket: WebSocket, user_id: str):
    await websocket.accept()
    local_connections[user_id] = websocket
    try:
        while True:
            data = await websocket.receive_text()
            # Publish to Redis — all instances (including this one) receive it
            await redis.publish(CHANNEL, f"{user_id}: {data}")
    except WebSocketDisconnect:
        local_connections.pop(user_id, None)

import asyncio
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from typing import Dict

app = FastAPI()


class HeartbeatManager:
    """Connection manager with per-connection heartbeat tasks.

    Each connection gets an independent heartbeat coroutine. If the client
    fails to respond to a ping within the timeout window, the connection is
    closed and removed from the active set.
    """

    def __init__(self, interval: int = 30, timeout: int = 10):
        self.connections: Dict[str, WebSocket] = {}
        self._last_pong: Dict[str, float] = {}
        self.interval = interval
        self.timeout = timeout

    async def connect(self, user_id: str, websocket: WebSocket) -> None:
        await websocket.accept()
        self.connections[user_id] = websocket
        self._last_pong[user_id] = asyncio.get_event_loop().time()
        asyncio.create_task(self._heartbeat_loop(user_id, websocket))

    def record_pong(self, user_id: str) -> None:
        """Called when the client sends a pong-type message."""
        self._last_pong[user_id] = asyncio.get_event_loop().time()

    def disconnect(self, user_id: str) -> None:
        self.connections.pop(user_id, None)
        self._last_pong.pop(user_id, None)

    async def _heartbeat_loop(self, user_id: str, ws: WebSocket) -> None:
        while user_id in self.connections:
            await asyncio.sleep(self.interval)
            if user_id not in self.connections:
                break
            # Check when we last heard from this client
            elapsed = asyncio.get_event_loop().time() - self._last_pong.get(user_id, 0)
            if elapsed > self.interval + self.timeout:
                # Client has not responded within the timeout window — treat as dead
                await ws.close(code=1001)  # 1001 = Going Away
                self.disconnect(user_id)
                break
            try:
                await ws.send_json({"type": "ping"})
            except Exception:
                # Send failed — connection is already dead at the network layer
                self.disconnect(user_id)
                break


manager = HeartbeatManager(interval=30, timeout=10)


@app.websocket("/ws/{user_id}")
async def ws_endpoint(websocket: WebSocket, user_id: str):
    await manager.connect(user_id, websocket)
    try:
        while True:
            data = await websocket.receive_json()
            if data.get("type") == "pong":
                manager.record_pong(user_id)
                continue  # Heartbeat response — no further processing needed
            # Handle actual application messages here
            await websocket.send_json({"echo": data})
    except WebSocketDisconnect:
        manager.disconnect(user_id)

import uvicorn

# Production WebSocket server configuration
# Run this with: python -m io.thecodeforge.config.uvicorn_ws
# Or directly: uvicorn app:app --workers 4 --loop uvloop

if __name__ == "__main__":
    uvicorn.run(
        "io.thecodeforge.realtime.connection_manager:app",
        host="0.0.0.0",
        port=8000,
        workers=4,                # One per physical CPU core for CPU-bound handler work
        loop="uvloop",            # 2-4x faster than default asyncio on I/O-bound workloads
        ws="websockets",          # websockets library — better performance than wsproto default
        ws_ping_interval=30,      # Send native WS ping frame every 30 seconds
        ws_ping_timeout=10,       # Close connection if no pong within 10 seconds
        limit_concurrency=10000,  # Reject new connections beyond this per worker
        limit_max_requests=50000, # Restart worker after N requests — guards against slow leaks
        timeout_keep_alive=5,     # HTTP keep-alive timeout (not WebSocket — separate setting)
        access_log=True,          # Log WS upgrade requests alongside HTTP — useful for debugging
    )