Skip to content
Codeloom
FastAPI

WebSocket Endpoints in FastAPI for Real-Time Features

Build real-time features with FastAPI WebSockets including chat rooms, live notifications, and connection management patterns.

·6 min read · By Codeloom
Intermediate 11 min read

What you'll learn

  • ✓Creating WebSocket endpoints in FastAPI
  • ✓Managing multiple client connections
  • ✓Building a real-time chat room with broadcast

Prerequisites

  • •Basic FastAPI knowledge
  • •Understanding of async/await in Python
  • •Familiarity with HTTP vs WebSocket protocols

Why WebSockets?

HTTP follows a request-response pattern: the client asks, the server answers. WebSockets open a persistent, bidirectional channel. Both sides can send messages at any time without polling. This makes WebSockets ideal for chat applications, live dashboards, notifications, collaborative editing, and real-time data feeds.

FastAPI has first-class WebSocket support built on Starlette.

A Basic WebSocket Endpoint

from fastapi import FastAPI, WebSocket, WebSocketDisconnect

app = FastAPI()

@app.websocket("/ws")
async def websocket_endpoint(ws: WebSocket):
    await ws.accept()
    try:
        while True:
            data = await ws.receive_text()
            await ws.send_text(f"Echo: {data}")
    except WebSocketDisconnect:
        print("Client disconnected")

The flow is straightforward. The client connects, the server calls await ws.accept() to complete the handshake, and then both sides exchange messages in a loop. When the client disconnects, WebSocketDisconnect is raised.

Testing with a Browser

# Add an HTML page for quick testing
from fastapi.responses import HTMLResponse

@app.get("/")
async def get():
    return HTMLResponse("""
    <html>
    <body>
        <h1>WebSocket Test</h1>
        <input id="msg" type="text" />
        <button onclick="send()">Send</button>
        <ul id="messages"></ul>
        <script>
            const ws = new WebSocket("ws://localhost:8000/ws");
            ws.onmessage = (event) => {
                const li = document.createElement("li");
                li.textContent = event.data;
                document.getElementById("messages").appendChild(li);
            };
            function send() {
                const input = document.getElementById("msg");
                ws.send(input.value);
                input.value = "";
            }
        </script>
    </body>
    </html>
    """)

Connection Manager for Multiple Clients

Real applications need to track active connections, broadcast messages, and handle disconnections cleanly.

from typing import Dict, List, Set
from fastapi import WebSocket

class ConnectionManager:
    def __init__(self):
        self.active_connections: Dict[str, Set[WebSocket]] = {}

    async def connect(self, websocket: WebSocket, room: str):
        await websocket.accept()
        if room not in self.active_connections:
            self.active_connections[room] = set()
        self.active_connections[room].add(websocket)

    def disconnect(self, websocket: WebSocket, room: str):
        if room in self.active_connections:
            self.active_connections[room].discard(websocket)
            if not self.active_connections[room]:
                del self.active_connections[room]

    async def send_personal(self, message: str, websocket: WebSocket):
        await websocket.send_text(message)

    async def broadcast(self, message: str, room: str, exclude: WebSocket = None):
        if room not in self.active_connections:
            return
        disconnected = []
        for connection in self.active_connections[room]:
            if connection == exclude:
                continue
            try:
                await connection.send_text(message)
            except Exception:
                disconnected.append(connection)
        for conn in disconnected:
            self.active_connections[room].discard(conn)

manager = ConnectionManager()

The manager tracks connections grouped by room. Broadcasting iterates over all connections in a room and handles stale connections gracefully by removing them on failure.

Building a Chat Room

import json
from datetime import datetime

@app.websocket("/ws/chat/{room_name}")
async def chat_endpoint(
    websocket: WebSocket,
    room_name: str,
    username: str = "anonymous"
):
    await manager.connect(websocket, room_name)

    # Announce the user joined
    join_msg = json.dumps({
        "type": "system",
        "message": f"{username} joined the room",
        "timestamp": datetime.utcnow().isoformat()
    })
    await manager.broadcast(join_msg, room_name)

    try:
        while True:
            data = await websocket.receive_text()

            # Parse and validate the message
            try:
                payload = json.loads(data)
                content = payload.get("message", "").strip()
            except json.JSONDecodeError:
                content = data.strip()

            if not content:
                continue

            # Broadcast the message to all users in the room
            chat_msg = json.dumps({
                "type": "chat",
                "username": username,
                "message": content,
                "timestamp": datetime.utcnow().isoformat()
            })
            await manager.broadcast(chat_msg, room_name)

    except WebSocketDisconnect:
        manager.disconnect(websocket, room_name)
        leave_msg = json.dumps({
            "type": "system",
            "message": f"{username} left the room",
            "timestamp": datetime.utcnow().isoformat()
        })
        await manager.broadcast(leave_msg, room_name)

Users connect to /ws/chat/general?username=alice. The endpoint handles join announcements, message broadcasting, and leave notifications.

Authentication on WebSocket Connections

WebSockets do not support custom headers the same way HTTP does. Common authentication strategies include query parameters, cookies, and first-message auth.

from fastapi import Query, status
import jwt

SECRET_KEY = "your-secret-key"

async def verify_ws_token(token: str) -> dict:
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=["HS256"])
        return payload
    except jwt.InvalidTokenError:
        return None

@app.websocket("/ws/secure")
async def secure_websocket(
    websocket: WebSocket,
    token: str = Query(...)
):
    # Validate token before accepting the connection
    user = await verify_ws_token(token)
    if not user:
        await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
        return

    await websocket.accept()
    await websocket.send_text(f"Welcome, {user['username']}")

    try:
        while True:
            data = await websocket.receive_text()
            await websocket.send_text(f"[{user['username']}]: {data}")
    except WebSocketDisconnect:
        pass

Validating the token before calling accept() rejects unauthorized connections before the WebSocket handshake completes.

Heartbeat and Keep-Alive

Long-lived connections can silently drop due to network issues, proxies, or load balancers. Implement ping/pong heartbeats to detect dead connections.

import asyncio

@app.websocket("/ws/heartbeat")
async def heartbeat_endpoint(websocket: WebSocket):
    await websocket.accept()

    async def send_ping():
        while True:
            try:
                await asyncio.sleep(30)
                await websocket.send_json({"type": "ping"})
            except Exception:
                break

    ping_task = asyncio.create_task(send_ping())

    try:
        while True:
            data = await websocket.receive_json()

            if data.get("type") == "pong":
                continue  # Heartbeat response, ignore

            # Handle actual messages
            await websocket.send_json({
                "type": "response",
                "data": data
            })
    except WebSocketDisconnect:
        ping_task.cancel()

The server sends a ping every 30 seconds. If the client does not respond, the next ping attempt will fail and break the loop, cleaning up the connection.

Broadcasting from HTTP Endpoints

Sometimes you need to push updates to WebSocket clients from a regular HTTP endpoint — for example, when an admin action should notify connected users.

@app.post("/notifications/broadcast")
async def broadcast_notification(message: str, room: str = "general"):
    notification = json.dumps({
        "type": "notification",
        "message": message,
        "timestamp": datetime.utcnow().isoformat()
    })
    await manager.broadcast(notification, room)
    count = len(manager.active_connections.get(room, set()))
    return {"sent_to": count, "room": room}

This lets your REST API trigger real-time updates to connected WebSocket clients.

Handling JSON and Binary Data

WebSockets support both text and binary frames. FastAPI provides methods for each.

@app.websocket("/ws/data")
async def data_endpoint(websocket: WebSocket):
    await websocket.accept()
    try:
        while True:
            # Receive and send JSON
            data = await websocket.receive_json()

            if data["action"] == "subscribe":
                await websocket.send_json({
                    "status": "subscribed",
                    "channel": data["channel"]
                })
            elif data["action"] == "upload":
                # Switch to binary for file data
                binary_data = await websocket.receive_bytes()
                size = len(binary_data)
                await websocket.send_json({
                    "status": "received",
                    "bytes": size
                })
    except WebSocketDisconnect:
        pass

Scaling WebSockets Beyond a Single Server

The ConnectionManager above stores connections in memory, which means it only works on a single server. For multiple servers behind a load balancer, you need a pub/sub layer.

import redis.asyncio as redis

class RedisConnectionManager(ConnectionManager):
    def __init__(self, redis_url: str = "redis://localhost:6379"):
        super().__init__()
        self.redis = redis.from_url(redis_url)
        self.pubsub = self.redis.pubsub()

    async def start_listener(self, room: str):
        await self.pubsub.subscribe(f"chat:{room}")
        async for message in self.pubsub.listen():
            if message["type"] == "message":
                await self._local_broadcast(
                    message["data"].decode(),
                    room
                )

    async def broadcast(self, message: str, room: str, exclude=None):
        # Publish to Redis -- all servers receive it
        await self.redis.publish(f"chat:{room}", message)

    async def _local_broadcast(self, message: str, room: str):
        # Only broadcast to connections on THIS server
        await super().broadcast(message, room)

Each server publishes messages to Redis. Each server subscribes to channels and broadcasts only to its local connections. This gives you horizontal scaling for WebSocket applications.

Key Takeaways

FastAPI makes WebSocket endpoints straightforward with the @app.websocket decorator. Use a ConnectionManager class to track active connections and handle broadcasting. Authenticate connections before accepting them, implement heartbeats for reliability, and use Redis pub/sub when you need to scale beyond a single server. For most applications, start with the in-memory manager and add Redis when you introduce multiple server instances.

Related articles