Build a WebSocket connection manager with JWT auth

Introduction

When you ship a real-time feature like chat or collaboration, every WebSocket connection is a long-lived, stateful resource that holds a file descriptor and a slot in your event loop for its entire lifetime. If your server accepts connections without tracking them — or fails to detect when clients silently vanish on a flaky mobile network — the process leaks descriptors until it crashes or starves real users out of the system. By the end of this lesson, you'll be able to implement a connection manager that authenticates each handshake with a JWT, tracks per-connection state in a lookup structure, runs a heartbeat loop to evict dead sockets, and tears down deterministically on disconnect.

Key Terminology

• ConnectionManager: A server-side coordinator that authenticates each WebSocket handshake, registers per-connection state in a lookup structure, and tears down sockets on disconnect or timeout.
• JWT (JSON Web Token): A signed token passed by the client (here, as a ?token=... query parameter) that the server cryptographically verifies at handshake time to establish the authenticated user_id and token expiry for the connection.
• Heartbeat loop: A per-connection asyncio task that periodically sends a ping and tracks last_seen so the server can evict sockets that have gone silent (e.g., clients vanishing on a flaky mobile network) instead of waiting for OS-level TCP keepalive.

Concepts

The two subsections below explain why a naive dict of sockets fails under real-world conditions — silent mobile drops, mid-session token expiry, multi-instance horizontal scaling — and then name the four design decisions (token transport, ping strategy, idempotent teardown, dual-dictionary index) that turn the ConnectionManager into a leak-free coordinator.

Why Connection Management Is Non-Trivial at Scale

In a single-instance prototype, you can store WebSocket references in a Python dict and broadcast by iterating over it. In production, three forces break that model. First, clients on mobile networks silently vanish—TCP FIN packets never arrive, so the server believes the socket is alive until the OS-level keepalive timer fires (often 2+ hours by default). Second, authentication tokens expire mid-session; a connection accepted with a valid JWT at 2:00 PM may be operating on an expired token by 3:00 PM. Third, horizontal scaling means a user's connections may span multiple FastAPI instances, so local state alone cannot answer "who is online?" These constraints drive the design toward heartbeat-based liveness detection, per-connection metadata that includes token expiry, and an external store (Redis) as the coordination backbone for presence tracking covered in another goal.

Key Design Decisions for Production

• Token-in-query-parameter vs. subprotocol headers: Browser WebSocket APIs do not support custom headers on the upgrade request. Passing the JWT as a query parameter (?token=...) is the pragmatic choice. The tradeoff is that tokens appear in server access logs; configure your reverse proxy to strip or redact the token query parameter from logged URLs.

• Application-level ping vs. protocol-level ping: WebSocket protocol defines PING/PONG control frames, but many load balancers (ALB, Cloudflare) intercept or inject their own. Using JSON-level {"type": "ping"} messages guarantees end-to-end liveness detection under your control, independent of infrastructure behavior.

• Idempotent disconnect: The disconnect() method can be called multiple times safely—once by the heartbeat loop detecting a timeout, and again by the finally block in the endpoint. The pop(..., None) pattern and task done() check prevent double-close errors that would otherwise surface as noisy RuntimeError exceptions in production logs.

• Per-connection vs. per-user state: The dual-dictionary design (_connections + _user_connections) trades memory for lookup speed. A user with three open tabs has three ConnectionState entries but a single set in the reverse index. Broadcasting to a user scans only their connections, not the entire registry—an important optimization when the connection count reaches thousands.

Code Walkthrough

Now that you've seen why naive socket tracking fails and which design decisions address each failure mode, the walkthrough turns those decisions into concrete code. It proceeds in three stages: a Mermaid sequence diagram that traces a single connection from the HTTP upgrade through JWT verification to the heartbeat loop, the full ConnectionManager implementation in services/connection_manager.py, and the FastAPI @router.websocket("/ws") endpoint that drives it.

Connection Lifecycle Architecture

The following diagram traces a WebSocket connection from the initial HTTP upgrade request through authentication, registration, the heartbeat loop, and eventual teardown. Each transition corresponds to a method on the ConnectionManager class.

This sequence diagram maps the WebSocket handshake and session lifecycle between a Client, FastAPI Endpoint, ConnectionManager, JWT Verifier, and Redis presence store. When a client sends a GET /ws?token=eyJ... upgrade request, ConnectionManager.connect() delegates token validation to decode_and_verify(), which returns a UserClaims object containing user_id and exp. On success, ZADD registers the user in Redis for presence tracking, and a 15-second PING/PONG heartbeat loop continuously updates last_seen to detect stale connections—critical for accurate participant counts in shared AI conversations.

Loading diagram...

The diagram shows that authentication happens before the 101 response is sent—the client never enters an authenticated session if the JWT is invalid. The heartbeat loop runs as a concurrent asyncio task per connection, independent of the message-handling loop. This separation ensures that a slow LLM response (covered in another goal) does not delay liveness detection.

Core ConnectionManager Implementation

The ConnectionManager class in services/connection_manager.py is the central coordinator for all active WebSocket connections. It exposes three primary methods: connect() for authentication and registration, disconnect() for cleanup, and _heartbeat_loop() for liveness monitoring. Internally, it maintains two dictionaries—_connections mapping connection IDs to ConnectionState dataclass instances, and _user_connections providing a reverse index from user IDs to sets of connection IDs, enabling efficient per-user broadcasting. The ConnectionState dataclass captures the WebSocket reference, authenticated user ID, token expiry timestamp, and the last-seen heartbeat time. The following implementation demonstrates the complete lifecycle: handshake-time JWT verification, dual-index registration, idempotent disconnect, and a per-connection heartbeat loop that terminates on either liveness timeout or token expiry.

Code snippetpython
1import asyncio
2import time
3import uuid
4from dataclasses import dataclass, field
5from typing import Dict, Optional, Set
6
7import jwt
8from fastapi import WebSocket, WebSocketDisconnect
9
10JWT_SECRET = "your-secret-key"  # In production, load from environment
11JWT_ALGORITHM = "HS256"
12HEARTBEAT_INTERVAL = 15  # seconds
13HEARTBEAT_TIMEOUT = 45   # miss 3 beats = dead
14
15@dataclass
16class ConnectionState:
17    conn_id: str
18    websocket: WebSocket
19    user_id: str
20    token_exp: float
21    last_seen: float = field(default_factory=time.time)
22    heartbeat_task: Optional[asyncio.Task] = field(default=None, repr=False)
23
24class ConnectionManager:
25    def __init__(self):
26        self._connections: Dict[str, ConnectionState] = {}
27        self._user_connections: Dict[str, Set[str]] = {}
28
29    async def connect(self, websocket: WebSocket) -> ConnectionState:
30        token = websocket.query_params.get("token")
31        if token is None:
32            await websocket.close(code=4400, reason="Missing token")
33            raise ValueError("No token provided")
34
35        try:
36            payload = jwt.decode(token, JWT_SECRET, algorithms=[JWT_ALGORITHM])
37        except jwt.ExpiredSignatureError:
38            await websocket.close(code=4401, reason="Token expired")
39            raise ValueError("Expired token")
40        except jwt.InvalidTokenError:
41            await websocket.close(code=4401, reason="Invalid token")
42            raise ValueError("Invalid token")
43
44        await websocket.accept()
45
46        conn_id = str(uuid.uuid4())
47        state = ConnectionState(
48            conn_id=conn_id,
49            websocket=websocket,
50            user_id=payload["sub"],
51            token_exp=payload["exp"],
52        )
53
54        self._connections[conn_id] = state
55        self._user_connections.setdefault(payload["sub"], set()).add(conn_id)
56
57        state.heartbeat_task = asyncio.create_task(
58            self._heartbeat_loop(conn_id)
59        )
60        return state
61
62    async def disconnect(self, conn_id: str) -> None:
63        state = self._connections.pop(conn_id, None)
64        if state is None:
65            return
66        if state.heartbeat_task and not state.heartbeat_task.done():
67            state.heartbeat_task.cancel()
68
69        user_conns = self._user_connections.get(state.user_id, set())
70        user_conns.discard(conn_id)
71        if not user_conns:
72            self._user_connections.pop(state.user_id, None)
73
74        try:
75            await state.websocket.close(code=1000)
76        except RuntimeError:
77            pass  # Already closed
78
79    async def _heartbeat_loop(self, conn_id: str) -> None:
80        while conn_id in self._connections:
81            state = self._connections[conn_id]
82            elapsed = time.time() - state.last_seen
83            if elapsed > HEARTBEAT_TIMEOUT:
84                await self.disconnect(conn_id)
85                break
86            if time.time() > state.token_exp:
87                await self.disconnect(conn_id)
88                break
89            try:
90                await state.websocket.send_json({"type": "ping", "ts": time.time()})
91                await asyncio.sleep(HEARTBEAT_INTERVAL)
92            except (WebSocketDisconnect, RuntimeError):
93                await self.disconnect(conn_id)
94                break
95
96    def record_pong(self, conn_id: str) -> None:
97        state = self._connections.get(conn_id)
98        if state is not None:
99            state.last_seen = time.time()
100
101    def get_user_connections(self, user_id: str) -> list[ConnectionState]:
102        conn_ids = self._user_connections.get(user_id, set())
103        return [self._connections[cid] for cid in conn_ids if cid in self._connections]

• Imports and constants (lines 1-13): Standard library modules cover async scheduling, monotonic timestamps, UUID generation, dataclass definitions, and type annotations. PyJWT verifies tokens, and FastAPI's WebSocket / WebSocketDisconnect types provide the transport abstraction. The heartbeat constants encode the liveness contract: a ping every 15 seconds, and a connection declared dead after 45 seconds of silence (three missed beats). In production, JWT_SECRET must be injected from environment variables or a secrets manager—never hardcoded.
• ConnectionState dataclass (lines 16-23): Bundles all per-connection metadata into a single object. The last_seen field defaults to the current timestamp via field(default_factory=time.time), ensuring fresh connections start with a valid heartbeat baseline. The heartbeat_task holds a reference to the asyncio.Task so it can be cancelled during cleanup.
• __init__ (lines 26-29): Initializes two lookup structures. _connections provides O(1) access by connection ID. _user_connections is the reverse index enabling "send to all of user X's connections" without scanning every connection—critical when a single user has multiple browser tabs open.
• connect() (lines 31-62): Extracts the JWT from the query parameter. If the token is missing, the connection is closed with a custom 4400 code before accepting the upgrade. Custom close codes in the 4000-4999 range are reserved for application use per RFC 6455, enabling clients to distinguish authentication failures from server errors. JWT decoding uses explicit algorithm pinning (algorithms=[JWT_ALGORITHM]) to prevent algorithm-switching attacks where an attacker submits a token signed with HMAC using the public RSA key. Each failure mode—expired versus structurally invalid—receives a distinct close code. Only after successful verification does the server call websocket.accept(), which sends the 101 Switching Protocols response. The connection ID is a UUID4 rather than an auto-incrementing integer, preventing enumeration attacks and collisions across restarts. Registration into both lookup dictionaries uses setdefault to lazily initialize the per-user set, avoiding KeyError on the first connection from any user. The heartbeat task is spawned immediately via asyncio.create_task(), beginning the liveness monitoring loop.
• disconnect() (lines 64-79): Idempotent—calling it with an already-removed conn_id returns silently due to the pop(..., None) pattern. The heartbeat task is cancelled to prevent it from running against a stale state. The reverse index cleanup removes the user entry entirely when their last connection drops, keeping the dictionary from accumulating empty sets over time. The final close() is wrapped in a try/except RuntimeError because the socket may already be closed by the heartbeat loop or the client itself.
• _heartbeat_loop() (lines 81-96): Runs as a persistent coroutine per connection. It checks two termination conditions on each cycle: heartbeat timeout (client unresponsive) and token expiry (authentication lapsed). The ping message is sent as a JSON object rather than a WebSocket protocol-level ping frame, giving application code full control over the format and allowing clients to include metadata in the pong response.
• record_pong() (lines 98-101): Called from the message-handling loop when a pong arrives. It updates last_seen to the current time, resetting the timeout window. This is a deliberate separation—the heartbeat loop sends pings, but the message loop receives pongs, avoiding coupling between the two coroutines.
• get_user_connections() (lines 103-105): Returns all active ConnectionState objects for a given user, filtering out any stale IDs that might exist briefly during concurrent disconnect operations. This method feeds into the broadcasting logic covered in another goal.

Integrating with the FastAPI Endpoint

The FastAPI WebSocket endpoint ties the ConnectionManager into the actual request handling loop. The endpoint function below demonstrates the standard pattern: authenticate via connect(), enter the receive loop, dispatch messages by type, handle disconnection in a finally block. The record_pong call is placed inside the message dispatch to update heartbeat state whenever the client responds. This endpoint function in routers/ws.py accepts the WebSocket connection, delegates authentication to the ConnectionManager.connect() method, and then enters an indefinite receive loop that dispatches incoming JSON messages based on their type field—routing pong responses to the heartbeat tracker and chat messages to the conversation handler (detailed in another goal).

1from fastapi import APIRouter, WebSocket, WebSocketDisconnect
2from services.connection_manager import ConnectionManager
3
4router = APIRouter()
5manager = ConnectionManager()
6
7@router.websocket("/ws")
8async def websocket_endpoint(websocket: WebSocket):
9    state = None
10    try:
11        state = await manager.connect(websocket)
12        while True:
13            data = await websocket.receive_json()
14            msg_type = data.get("type", "")
15
16            if msg_type == "pong":
17                manager.record_pong(state.conn_id)
18            elif msg_type == "chat":
19                # Delegate to conversation handler
20                await handle_chat(state, data)
21            elif msg_type == "typing":
22                # Delegate to presence tracker
23                await handle_typing(state, data)
24            else:
25                await websocket.send_json(
26                    {"type": "error", "detail": f"Unknown type: {msg_type}"}
27                )
28    except WebSocketDisconnect:
29        pass
30    except ValueError:
31        return  # Auth failed, already closed
32    finally:
33        if state is not None:
34            await manager.disconnect(state.conn_id)

Do's and Don'ts

Do's

1. ✓Do authenticate the JWT via decode_and_verify() inside connect() before sending the 101 Switching Protocols response — closing with code 4401 on a ValueError ensures unauthenticated clients never enter an active session, so no file descriptor or event-loop slot is allocated to an unauthorized socket.
2. ✓Do maintain both _connections (connection ID → ConnectionState) and the _user_connections reverse index (user ID → set of connection IDs) — the dual-index structure lets you broadcast to every socket belonging to one user without scanning the entire _connections dict, which degrades linearly as active connections grow.
3. ✓Do spawn _heartbeat_loop() as a separate asyncio task per connection, sending a PING every 15 seconds and updating last_seen on each PONG — decoupling liveness detection from the message-handling coroutine means a slow LLM response cannot delay eviction of dead sockets, keeping participant presence counts in Redis accurate.

Don'ts

1. ✗Don't defer JWT validation until after the WebSocket upgrade completes — accepting the 101 before calling decode_and_verify() lets unauthenticated clients hold a live socket and a file descriptor indefinitely, bypassing the 4401 close path and leaking resources until the process exhausts its descriptor limit.
2. ✗Don't track connections with only a single _connections dict keyed by connection ID — without _user_connections as a reverse index, sending a message to all sockets owned by a specific user requires iterating every active ConnectionState and filtering by user_id, an O(N) scan that becomes expensive in a shared AI conversation with many concurrent participants.
3. ✗Don't merge the _heartbeat_loop() logic into the same coroutine that reads incoming messages — if that coroutine blocks waiting on an upstream LLM response, PING frames are delayed past the 15-second window, causing live connections to time out incorrectly or dead sockets to persist undetected and stale in the Redis presence sorted set.