FastAPI WebSocket Routes

In FastAPI, WebSocket routes allow you to establish persistent, bidirectional communication channels between clients and your server. Unlike traditional HTTP endpoints that follow a request-response pattern, WebSockets enable continuous data exchange, making them perfect for applications requiring real-time updates like chat applications, live dashboards, or online games.

Understanding WebSocket Routing

WebSocket routes in FastAPI are defined similarly to regular HTTP endpoints but use a different decorator and handling mechanism to maintain the long-lived connection.

Basic WebSocket Route

Let's start with a simple WebSocket route:

python

from fastapi import FastAPI, WebSocket

app = FastAPI()

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    while True:
        data = await websocket.receive_text()
        await websocket.send_text(f"Message received: {data}")

In this basic example:

1. We use the @app.websocket decorator to define a WebSocket endpoint at the "/ws" path
2. The function receives a WebSocket instance as a parameter
3. We accept the connection with await websocket.accept()
4. We enter a loop that continuously receives and sends messages

WebSocket Route Paths

Like REST endpoints, WebSocket routes can have:

Fixed Paths

python

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    # Handle connection

Path Parameters

Path parameters allow dynamic values in your WebSocket routes:

python

@app.websocket("/ws/{client_id}")
async def websocket_endpoint(websocket: WebSocket, client_id: str):
    await websocket.accept()
    await websocket.send_text(f"Connected to client: {client_id}")
    while True:
        data = await websocket.receive_text()
        await websocket.send_text(f"Client {client_id} sent: {data}")

This is particularly useful for identifying different clients or channels in applications like chat rooms or user-specific dashboards.

WebSocket Connection Managers

For more complex applications, especially those with multiple clients connecting simultaneously, it's helpful to create a connection manager class:

python

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from typing import List

app = FastAPI()

class ConnectionManager:
    def __init__(self):
        self.active_connections: List[WebSocket] = []
    
    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)
    
    def disconnect(self, websocket: WebSocket):
        self.active_connections.remove(websocket)
    
    async def broadcast(self, message: str):
        for connection in self.active_connections:
            await connection.send_text(message)

manager = ConnectionManager()

@app.websocket("/ws/{client_id}")
async def websocket_endpoint(websocket: WebSocket, client_id: str):
    await manager.connect(websocket)
    try:
        while True:
            data = await websocket.receive_text()
            # You can process the data here
            await manager.broadcast(f"Client #{client_id} says: {data}")
    except WebSocketDisconnect:
        manager.disconnect(websocket)
        await manager.broadcast(f"Client #{client_id} left the chat")

This connection manager pattern allows you to:

1. Track all active connections
2. Send messages to all connected clients (broadcasting)
3. Handle client disconnections gracefully

Organizing WebSocket Routes with APIRouter

For larger applications, you can organize WebSocket routes using FastAPI's APIRouter, just like with regular HTTP routes:

python

from fastapi import APIRouter, WebSocket

router = APIRouter()

@router.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    while True:
        data = await websocket.receive_text()
        await websocket.send_text(f"You sent: {data}")

# In your main app file
from fastapi import FastAPI
from .routes import router

app = FastAPI()
app.include_router(router, prefix="/chat")

Now your WebSocket endpoint will be available at "/chat/ws".

Real-World Example: Chat Application

Let's build a simple chat application to demonstrate WebSocket routes in action:

python

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse
from typing import List, Dict

app = FastAPI()

# HTML for a basic chat client
html = """
<!DOCTYPE html>
<html>
    <head>
        <title>FastAPI Chat</title>
    </head>
    <body>
        <h1>FastAPI Chat</h1>
        <h2>Your ID: <span id="ws-id"></span></h2>
        <input type="text" id="messageText" placeholder="Type a message"/>
        <button onclick="sendMessage()">Send</button>
        <ul id="messages">
        </ul>
        <script>
            var client_id = Date.now()
            document.querySelector("#ws-id").textContent = client_id;
            var ws = new WebSocket(`ws://localhost:8000/ws/${client_id}`);
            ws.onmessage = function(event) {
                var messages = document.getElementById('messages')
                var message = document.createElement('li')
                var content = document.createTextNode(event.data)
                message.appendChild(content)
                messages.appendChild(message)
            };
            function sendMessage() {
                var messageInput = document.getElementById('messageText')
                ws.send(messageInput.value)
                messageInput.value = ''
            }
        </script>
    </body>
</html>
"""

class ChatConnectionManager:
    def __init__(self):
        self.active_connections: Dict[str, WebSocket] = {}
    
    async def connect(self, client_id: str, websocket: WebSocket):
        await websocket.accept()
        self.active_connections[client_id] = websocket
    
    def disconnect(self, client_id: str):
        if client_id in self.active_connections:
            del self.active_connections[client_id]
    
    async def send_personal_message(self, message: str, client_id: str):
        if client_id in self.active_connections:
            await self.active_connections[client_id].send_text(message)
    
    async def broadcast(self, message: str, exclude_client: str = None):
        for client_id, connection in self.active_connections.items():
            if client_id != exclude_client:
                await connection.send_text(message)

manager = ChatConnectionManager()

@app.get("/")
async def get():
    return HTMLResponse(html)

@app.websocket("/ws/{client_id}")
async def websocket_endpoint(websocket: WebSocket, client_id: str):
    await manager.connect(client_id, websocket)
    try:
        await manager.broadcast(f"Client #{client_id} joined the chat")
        while True:
            data = await websocket.receive_text()
            await manager.broadcast(f"Client #{client_id}: {data}", exclude_client=client_id)
            await manager.send_personal_message(f"You: {data}", client_id)
    except WebSocketDisconnect:
        manager.disconnect(client_id)
        await manager.broadcast(f"Client #{client_id} left the chat")

This example demonstrates:

1. A more advanced connection manager that maps client IDs to their WebSocket connections
2. Broadcasting messages to all clients except the sender
3. Sending personal messages to specific clients
4. Providing a simple HTML interface that connects to our WebSocket endpoint

Best Practices for WebSocket Routes

1. Error Handling: Always wrap your WebSocket handling code in try/except blocks to gracefully handle disconnections.

python

try:
    while True:
        data = await websocket.receive_text()
        # Process data
except WebSocketDisconnect:
    # Clean up resources

2. Connection Timeouts: Implement timeout mechanisms for inactive connections:

python

import asyncio

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    try:
        while True:
            # Set a timeout of 60 seconds
            data = await asyncio.wait_for(
                websocket.receive_text(), 
                timeout=60.0
            )
            await websocket.send_text(f"You sent: {data}")
    except asyncio.TimeoutError:
        await websocket.close(code=1000, reason="Session timeout")
    except WebSocketDisconnect:
        # Handle disconnect

3. Authentication: Protect your WebSocket routes with authentication:

python

from fastapi import Depends, Cookie, Query

async def get_user(
    token: str = Query(None),
    session: str = Cookie(None)
):
    if token is None and session is None:
        return None
    # Validate token or session
    # Return user or None

@app.websocket("/ws")
async def websocket_endpoint(
    websocket: WebSocket,
    user: dict = Depends(get_user)
):
    if user is None:
        await websocket.close(code=1008, reason="Not authenticated")
        return
    
    await websocket.accept()
    # Proceed with authenticated user

Summary

WebSocket routes in FastAPI provide a powerful way to implement real-time communication features in your applications. By understanding how to:

• Create basic WebSocket endpoints
• Use path parameters for dynamic routing
• Implement connection managers for handling multiple clients
• Organize routes with APIRouter
• Apply authentication and error handling

You can build robust real-time applications like chat systems, live dashboards, collaborative tools, and more.

Exercises

1. Build a simple echo server that returns whatever message the client sends
2. Create a WebSocket-based notification system that sends periodic updates to all connected clients
3. Implement a chat room system where users can join different rooms based on room IDs
4. Add authentication to your WebSocket routes using JWT tokens
5. Create a live collaborative drawing application where multiple users can draw on the same canvas

Additional Resources

• FastAPI WebSockets Documentation
• MDN WebSocket API Documentation
• WebSockets Protocol RFC6455
• Socket.IO - Another popular library for real-time web applications

Happy coding with FastAPI WebSockets!