The core of any real-time chat application is its ability to establish persistent, bidirectional communication channels. This is where WebSockets come in. In this chapter, we’ll integrate a basic WebSocket endpoint into our FastAPI application.

Purpose of this Chapter

By the end of this chapter, you will be able to:

  • Understand the fundamental concept of WebSockets.
  • Implement a FastAPI WebSocket endpoint.
  • Send and receive messages over a WebSocket connection.
  • Test your WebSocket connection using a simple client.

Concepts Explained: WebSockets

Traditionally, HTTP is a request-response protocol: a client sends a request, the server sends a response, and the connection closes. This isn’t ideal for real-time applications where information needs to be pushed from the server to the client instantly, and vice-versa, without constant polling.

WebSockets provide a persistent, full-duplex communication channel over a single TCP connection. Once established (after an initial HTTP handshake), both the client and server can send data to each other at any time, making it perfect for real-time features like chat, live updates, and gaming.

FastAPI makes it straightforward to define WebSocket endpoints using the @app.websocket() decorator.

Step-by-Step Tasks

1. Add a WebSocket Endpoint to app/main.py

Let’s modify our app/main.py to include a WebSocket endpoint.

# app/main.py

from fastapi import FastAPI, WebSocket, WebSocketDisconnect

app = FastAPI()

@app.get("/")
async def read_root():
    return {"message": "Welcome to the Real-time Chat API!"}

@app.get("/health")
async def health_check():
    return {"status": "ok"}

# Our first WebSocket endpoint
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept() # Accept the WebSocket connection
    try:
        while True:
            # Receive text data from the client
            data = await websocket.receive_text()
            print(f"Received message: {data}")
            # Send the received data back to the client
            await websocket.send_text(f"Message text was: {data}")
    except WebSocketDisconnect:
        print("Client disconnected.")

Code Explanation:

  • from fastapi import FastAPI, WebSocket, WebSocketDisconnect: We import WebSocket to handle WebSocket connections and WebSocketDisconnect for exception handling.
  • @app.websocket("/ws"): This decorator registers the websocket_endpoint function to handle WebSocket connections at the /ws path.
  • async def websocket_endpoint(websocket: WebSocket): The function takes a WebSocket object as an argument, representing the active connection.
  • await websocket.accept(): This is crucial! Before you can send or receive data, the server must explicitly accept the incoming WebSocket connection.
  • while True:: We enter an infinite loop to continuously listen for messages from the client.
  • data = await websocket.receive_text(): Asynchronously waits to receive a text message from the client.
  • print(f"Received message: {data}"): Prints the received message to the server console.
  • await websocket.send_text(f"Message text was: {data}"): Echoes the received message back to the client, prefixing it.
  • except WebSocketDisconnect:: This block gracefully handles when a client closes the WebSocket connection.

2. Run the FastAPI Application

Ensure your virtual environment is activated and run Uvicorn:

pipenv shell
uvicorn app.main:app --reload

3. Test with a Simple JavaScript Client (HTML)

To test our WebSocket, we need a client. We’ll use a basic HTML page with JavaScript. Create a file named client.html outside your realtime-chat-app directory, or in a temporary location, for now.

<!-- client.html -->
<!DOCTYPE html>
<html>
<head>
    <title>FastAPI WebSocket Chat Client</title>
</head>
<body>
    <h1>WebSocket Chat Test</h1>
    <input type="text" id="messageInput" placeholder="Type a message">
    <button onclick="sendMessage()">Send</button>
    <div id="messages"></div>

    <script>
        const ws = new WebSocket("ws://localhost:8000/ws");
        const messagesDiv = document.getElementById("messages");
        const messageInput = document.getElementById("messageInput");

        ws.onopen = (event) => {
            messagesDiv.innerHTML += "<p>Connected to WebSocket server!</p>";
            console.log("WebSocket opened:", event);
        };

        ws.onmessage = (event) => {
            messagesDiv.innerHTML += `<p>Received: ${event.data}</p>`;
            console.log("WebSocket message:", event.data);
        };

        ws.onclose = (event) => {
            messagesDiv.innerHTML += "<p>Disconnected from WebSocket server.</p>";
            console.log("WebSocket closed:", event);
        };

        ws.onerror = (event) => {
            messagesDiv.innerHTML += "<p style='color:red;'>WebSocket error!</p>";
            console.error("WebSocket error:", event);
        };

        function sendMessage() {
            const message = messageInput.value;
            if (message) {
                ws.send(message);
                messagesDiv.innerHTML += `<p>Sent: ${message}</p>`;
                messageInput.value = ""; // Clear input field
            }
        }

        // Send a message when Enter key is pressed
        messageInput.addEventListener("keypress", (event) => {
            if (event.key === "Enter") {
                sendMessage();
            }
        });
    </script>
</body>
</html>

Open this client.html file in your web browser. Make sure your FastAPI server is running.

You should see “Connected to WebSocket server!”. Type messages into the input field and click “Send” or press Enter. You’ll see your messages echoed back, prefixed with “Message text was:”. Check your server’s terminal, and you’ll see the “Received message:” output.

If you close the browser tab, the server will print “Client disconnected.”

Tips/Challenges/Errors

  • WebSocketHandshakeError / “WebSocket connection failed”: This often means the server wasn’t running when the client tried to connect, or the WebSocket URL (ws://localhost:8000/ws) is incorrect.
  • Browser Security (Mixed Content): If your main website is HTTPS (secure) and your WebSocket is WS (insecure), browsers might block the connection. For local development, http and ws are fine. In production, you’ll typically use https for your main app and wss for WebSockets.
  • Server not printing “Client disconnected”: This might happen if the server crashes instead of the client gracefully closing the connection. Ensure your try...except WebSocketDisconnect block is correct.

Summary/Key Takeaways

You’ve successfully implemented a fundamental WebSocket endpoint in FastAPI, allowing for real-time, bidirectional communication between your server and a client. This is the bedrock of our chat application. However, a real chat application needs to manage multiple connections and broadcast messages. We’ll tackle this in the next chapter.