WebSockets — Module 18.2

How This Lesson Fits the Module & Volume

Celery jobs leave the client polling GET /v1/jobs/{id}. That works; it feels laggy. WebSockets keep a bidirectional channel so the server can push job status, collaborative agent traces, or voice-turn signals. The next lecture—Streaming—covers token-by-token LLM output, often via SSE instead of WS. Learn both; pick on purpose.

FastAPI/Starlette supports WebSockets natively. Flask needs extra ASGI help. Auth still applies: a socket is not a backdoor around authentication.

Learning Objectives

By the end of this lesson, students should be able to:

Describe the WS handshake (HTTP Upgrade) vs ordinary REST.
Push Celery job events to a connected client.
Contrast WebSockets vs SSE vs polling vs sync POST.
Authenticate the upgrade (query token or first-message JWT)—never skip.
Plan fan-out with Redis pub/sub when many API replicas exist.
Know when WS is overkill for one-way token streams.

Definition

A WebSocket is a persistent, full-duplex TCP connection initiated with an HTTP Upgrade: websocket handshake, then framed messages (text/binary) in both directions. Unlike HTTP request/response, either peer can send at any time until close. It is a transport, not an application protocol—you still define JSON message types (job.progress, chat.user, …).

Choose a Delivery Mode

Mode	Direction	Typical AI use
Sync HTTP	Req → res	Short chat, embeddings
Polling	Client pulls	Simple job status
SSE	Server → client	Token streaming (next lecture)
WebSocket	Both ways	Jobs + user interrupts + collab + voice

FastAPI WS + Job Events

# pip install fastapi uvicorn redis
# Client: new WebSocket("wss://api.example/v1/ws?token=...")
import asyncio, json, os
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Query, status
import redis.asyncio as redis

app = FastAPI()
r = redis.from_url(os.environ["REDIS_URL"])

def token_ok(token: str | None) -> bool:
    return bool(token) and token == os.environ.get("DEMO_WS_TOKEN", "")

@app.websocket("/v1/ws")
async def ws_jobs(ws: WebSocket, token: str | None = Query(default=None)):
    if not token_ok(token):
        await ws.close(code=status.WS_1008_POLICY_VIOLATION)
        return
    await ws.accept()
    pubsub = r.pubsub()
    await pubsub.subscribe("jobs")  # workers PUBLISH here on progress/done
    try:
        while True:
            recv = asyncio.create_task(ws.receive_text())
            msg = asyncio.create_task(pubsub.get_message(ignore_subscribe_messages=True, timeout=1.0))
            done, pending = await asyncio.wait({recv, msg}, return_when=asyncio.FIRST_COMPLETED)
            for t in pending:
                t.cancel()
            if recv in done:
                user_msg = json.loads(recv.result())
                # e.g. {"type": "cancel", "job_id": "..."} — bidirectional
                await ws.send_json({"type": "ack", "echo": user_msg.get("type")})
            if msg in done:
                m = msg.result()
                if m and m.get("data"):
                    await ws.send_text(m["data"] if isinstance(m["data"], str) else m["data"].decode())
    except WebSocketDisconnect:
        await pubsub.unsubscribe("jobs")

# Worker (Celery): r.publish("jobs", json.dumps({"type":"job.done","id": job_id}))

Multi-Replica Reality

One Uvicorn

In-memory client set works
Fine for class demos
Breaks when you scale Pods

Many API Pods

Redis pub/sub or Streams
Sticky sessions optional
User may land on any replica

Prefer SSE when

Only server → client tokens
HTTP/2 + proxies friendlier
Simpler auth (same headers)

WS strengths

Cancel mid-job, barge-in voice
Low overhead vs naive polling
Agent traces as event streams

WS costs

Load balancers must support Upgrade
Idle timeouts, heartbeats
Auth + backpressure are on you

Related Lectures

Lecture	Role
Celery	Produces the events you push
Streaming	SSE/token pattern next
Redis	Pub/sub across replicas
Authentication	Upgrade is still an authn event

Common Misconception

“WebSockets are automatically authenticated because the user loaded our SPA.” Cookies may not be sent the way you think; validate a token on accept. Second: WS is not required for ChatGPT-style typing—SSE is often enough. Third: opening a socket to OpenAI from the browser still leaks keys. Fourth: forgetting heartbeats + LB idle timeout looks like “random disconnects.”

Knowledge Check

Short Answer: How does a WebSocket start? Answer: HTTP Upgrade handshake, then framed bidirectional messages.
True/False: WebSockets are full-duplex. Answer: True.
Multiple Choice: Best first choice for one-way token stream: (a) SSE, (b) SMTP, (c) k-means. Answer: (a).
Short Answer: Why Redis pub/sub with many API replicas? Answer: The worker’s event must reach whichever Pod holds the user’s socket.
True/False: You can skip auth on /v1/ws if REST is authenticated. Answer: False.
Multiple Choice: Rejecting a bad WS token often uses close code: (a) 1008, (b) 204, (c) CFG. Answer: (a) (policy violation; teaching sketch).
Short Answer: Name a bidirectional AI UX that SSE cannot do as naturally. Answer: Any of: cancel/barge-in, collaborative edits, client audio chunks up + tokens down.
True/False: Polling job status is always wrong. Answer: False—simple and proxy-friendly; WS is an upgrade.
Multiple Choice: Next lecture: (a) Streaming, (b) DreamBooth, (c) Ridge. Answer: (a).
Short Answer: Why do LBs matter for WS? Answer: They must allow Upgrade and long-lived connections / idle timeouts / heartbeats.

Key Takeaways

WebSockets = bidirectional, persistent transport after HTTP Upgrade.
Push Celery events; auth the handshake; fan-out via Redis when scaled.
SSE often wins for one-way tokens.
Not a substitute for a product REST API.
Next: Streaming.

Trainer’s Guide

Lab: Browser mini-client: connect WS, enqueue Celery job via REST, receive job.done. Then implement cancel over the socket.

Whiteboard: Polling vs WS vs SSE. Circle auth on each. Arrow Redis between worker and two API Pods.

Recap: WebSockets push live events both ways. Token UX next: Streaming.