Skip to main content

FastAPI WebSockets Introduction

What are WebSockets?​

WebSockets provide a persistent connection between a client and server that both parties can use to start sending data at any time. Unlike the traditional HTTP request-response model, WebSockets enable real-time communication where both the client and server can send messages independently without needing to establish a new connection each time.

Traditional HTTP connections:

  • Client makes a request
  • Server sends a response
  • Connection closes

WebSocket connections:

  • Client makes an initial handshake request
  • Connection is upgraded from HTTP to WebSocket
  • Connection remains open for continuous two-way communication
  • Either party can send messages at any time

Why Use WebSockets with FastAPI?​

FastAPI's asynchronous nature makes it perfect for implementing WebSockets:

  • Real-time applications: Chat applications, live dashboards, collaborative tools
  • Reduced latency: No need to establish new connections for each message
  • Server-initiated updates: Server can push updates to clients without polling
  • Efficient resource usage: Fewer connections needed compared to HTTP polling

Basic WebSocket Implementation in FastAPI​

Let's start with a simple WebSocket example. First, ensure you have the necessary dependencies:

bash
pip install fastapi uvicorn websockets

Now, let's create a basic echo server that receives messages and sends them back:

python
from fastapi import FastAPI, WebSocket
import uvicorn

app = FastAPI()

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    try:
        while True:
            data = await websocket.receive_text()
            await websocket.send_text(f"Message received: {data}")
    except Exception as e:
        print(f"Error: {e}")
    finally:
        await websocket.close()

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)

Let's break down this code:

  1. We import the necessary modules: FastAPI and WebSocket
  2. We create a FastAPI application instance
  3. We define an endpoint with the @app.websocket decorator
  4. Inside the endpoint function:
    • We accept the WebSocket connection with await websocket.accept()
    • We enter a loop that continuously receives and sends messages
    • We handle any exceptions and ensure the connection closes properly

Testing Your WebSocket​

You can test your WebSocket implementation using a variety of tools:

Browser-based Testing​

Create a simple HTML file named websocket_test.html:

html
<!DOCTYPE html>
<html>
<head>
    <title>FastAPI WebSocket Test</title>
</head>
<body>
    <h1>WebSocket Test</h1>
    <form action="" onsubmit="sendMessage(event)">
        <input type="text" id="messageText" autocomplete="off"/>
        <button>Send</button>
    </form>
    <ul id="messages">
    </ul>
    <script>
        var ws = new WebSocket("ws://localhost:8000/ws");
        
        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(event) {
            var input = document.getElementById("messageText");
            ws.send(input.value);
            input.value = '';
            event.preventDefault();
        }
    </script>
</body>
</html>

Open this HTML file in your browser while your FastAPI application is running, and you'll be able to send messages and see the responses.

WebSocket Connection States​

WebSocket connections in FastAPI can be in several states:

  1. Connecting: Initial state when the connection request is made
  2. Connected: After websocket.accept() is called
  3. Disconnected: After websocket.close() or when an error occurs

It's important to properly manage these states to avoid resources remaining allocated when connections are no longer active.

Handling Multiple Clients​

For most real-world applications, you'll need to handle multiple WebSocket connections simultaneously. Let's implement a simple chat room:

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 send_personal_message(self, message: str, websocket: WebSocket):
        await websocket.send_text(message)

    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()
            # Send message to the sender
            await manager.send_personal_message(f"You: {data}", websocket)
            # Broadcast message to all other connected clients
            await manager.broadcast(f"Client #{client_id} says: {data}")
    except WebSocketDisconnect:
        manager.disconnect(websocket)
        await manager.broadcast(f"Client #{client_id} left the chat")

In this example:

  1. We create a ConnectionManager class to handle multiple WebSocket connections
  2. The manager keeps track of active connections and provides methods to send messages
  3. When a client connects, we add them to the list of active connections
  4. When a client disconnects, we remove them from the list
  5. We can send personal messages to specific clients or broadcast to all connected clients

Sending Different Types of Data​

WebSockets can transmit different types of data:

Text Data​

python
# Sending text
await websocket.send_text("Hello, client!")

# Receiving text
text_data = await websocket.receive_text()

JSON Data​

python
# Sending JSON
await websocket.send_json({"message": "Hello", "user": "system"})

# Receiving JSON
json_data = await websocket.receive_json()

Binary Data​

python
# Sending binary data
await websocket.send_bytes(b"Binary data here")

# Receiving binary data
binary_data = await websocket.receive_bytes()

Real-world Example: Live Dashboard​

Let's create a more practical example - a simple live dashboard that sends periodic updates:

python
import asyncio
import random
from fastapi import FastAPI, WebSocket, WebSocketDisconnect

app = FastAPI()

class DashboardManager:
    def __init__(self):
        self.active_connections = []
        self.generator_task = None

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)
        if self.generator_task is None:
            self.generator_task = asyncio.create_task(self.data_generator())

    def disconnect(self, websocket: WebSocket):
        self.active_connections.remove(websocket)
        if len(self.active_connections) == 0 and self.generator_task:
            self.generator_task.cancel()
            self.generator_task = None

    async def data_generator(self):
        while True:
            # Generate random metrics
            data = {
                "cpu_usage": random.randint(0, 100),
                "memory_usage": random.randint(0, 100),
                "active_users": random.randint(10, 1000),
                "response_time": round(random.uniform(0.1, 2.0), 2)
            }
            
            # Send to all connected clients
            for connection in self.active_connections:
                await connection.send_json(data)
            
            # Wait before sending the next update
            await asyncio.sleep(2)

dashboard = DashboardManager()

@app.websocket("/dashboard-feed")
async def dashboard_endpoint(websocket: WebSocket):
    await dashboard.connect(websocket)
    try:
        while True:
            # Just keep the connection alive and wait for data
            data = await websocket.receive_text()
    except WebSocketDisconnect:
        dashboard.disconnect(websocket)

This example:

  1. Creates a persistent WebSocket connection for dashboard clients
  2. Generates random system metrics every 2 seconds
  3. Broadcasts the metrics to all connected clients
  4. Efficiently manages resources by only generating data when clients are connected

You could create a simple frontend to display this data:

html
<!DOCTYPE html>
<html>
<head>
    <title>Live Dashboard</title>
    <style>
        .metric {
            padding: 10px;
            margin: 5px;
            border: 1px solid #ddd;
            border-radius: 5px;
            display: inline-block;
            width: 150px;
            text-align: center;
        }
        .value {
            font-size: 24px;
            font-weight: bold;
        }
    </style>
</head>
<body>
    <h1>System Dashboard</h1>
    <div id="dashboard">
        <div class="metric">
            <div>CPU Usage</div>
            <div id="cpu" class="value">0%</div>
        </div>
        <div class="metric">
            <div>Memory Usage</div>
            <div id="memory" class="value">0%</div>
        </div>
        <div class="metric">
            <div>Active Users</div>
            <div id="users" class="value">0</div>
        </div>
        <div class="metric">
            <div>Response Time</div>
            <div id="response" class="value">0 ms</div>
        </div>
    </div>

    <script>
        const ws = new WebSocket("ws://localhost:8000/dashboard-feed");
        
        ws.onmessage = function(event) {
            const data = JSON.parse(event.data);
            document.getElementById('cpu').textContent = data.cpu_usage + '%';
            document.getElementById('memory').textContent = data.memory_usage + '%';
            document.getElementById('users').textContent = data.active_users;
            document.getElementById('response').textContent = data.response_time + ' ms';
        };
        
        ws.onclose = function(event) {
            console.log("Connection closed");
        };
    </script>
</body>
</html>

Best Practices for WebSockets in FastAPI​

  1. Error Handling: Always implement proper error handling to prevent your application from crashing.
python
try:
    # WebSocket operations
except WebSocketDisconnect:
    # Handle normal disconnection
except Exception as e:
    # Handle unexpected errors
    logger.error(f"WebSocket error: {e}")
finally:
    # Clean up resources
  1. Authentication: Secure your WebSockets:
python
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    # Verify the token from the query parameters
    token = websocket.query_params.get("token")
    if not is_valid_token(token):
        await websocket.close(code=1008)  # Policy violation
        return
    
    await websocket.accept()
    # Proceed with authenticated connection

  1. Rate limiting: Prevent abuse by implementing rate limiting.

  2. Connection timeouts: Implement heartbeats to detect and clean up stale connections.

  3. Graceful degradation: Provide fallback mechanisms when WebSockets aren't supported.

Summary​

WebSockets in FastAPI enable real-time bidirectional communication between clients and servers. They're perfect for applications requiring live updates, chat functionality, collaborative editing, and more. FastAPI's asynchronous foundation makes it particularly well-suited for efficient WebSocket implementation.

In this introduction, we've covered:

  • The basics of WebSocket connections
  • Simple echo server implementation
  • Managing multiple client connections
  • Handling various data types
  • A practical dashboard example
  • Best practices for production use

Further Resources​

Exercises​

  1. Modify the chat example to display a list of active users.
  2. Create a collaborative drawing application where multiple users can draw on the same canvas.
  3. Implement a notification system that sends real-time alerts to admin users.
  4. Add secure authentication to the WebSocket connection using JWT tokens.
  5. Build a real-time multiplayer game using WebSockets for communication.

Happy coding with FastAPI WebSockets!


If you spot any mistakes on this website, please let me know at [email protected]. I’d greatly appreciate your feedback! :)