WebSockets

pip install "uvicorn[standard]"   # or: pip install websockets

from fastapi import FastAPI, WebSocket
from fastapi.testclient import TestClient

app = FastAPI()


@app.websocket("/ws/echo")
async def echo(websocket: WebSocket) -> None:
    await websocket.accept()
    async for text in websocket.iter_text():  # ends cleanly when the client disconnects
        await websocket.send_text(f"You said: {text}")
    print("client left")


client = TestClient(app)
with client.websocket_connect("/ws/echo") as ws:
    ws.send_text("hello")
    print(ws.receive_text())
    ws.send_text("again")
    print(ws.receive_text())

You said: hello
You said: again
client left

from fastapi import APIRouter, FastAPI, WebSocket
from fastapi.testclient import TestClient

router = APIRouter(prefix="/ws/books", tags=["chat"])


@router.websocket("/{book_id}/state")
async def state(websocket: WebSocket, book_id: int) -> None:
    before = websocket.application_state
    await websocket.accept(subprotocol="bookshelf.v1" if "bookshelf.v1" in websocket.scope["subprotocols"] else None)
    await websocket.send_json({
        "before": before.name, "after": websocket.application_state.name,
        "client": websocket.client_state.name, "offered": websocket.scope["subprotocols"],
        "path_params": websocket.path_params, "url": websocket.url.path,
    })
    await websocket.close()


app = FastAPI()
app.include_router(router)
client = TestClient(app)
with client.websocket_connect("/ws/books/3/state", subprotocols=["bookshelf.v1", "other"]) as ws:
    print(ws.accepted_subprotocol)
    print(ws.receive_json())
print(sorted(client.get("/openapi.json").json()["paths"]))

bookshelf.v1
{'before': 'CONNECTING', 'after': 'CONNECTED', 'client': 'CONNECTED', 'offered': ['bookshelf.v1', 'other'], 'path_params': {'book_id': '3'}, 'url': '/ws/books/3/state'}
[]

from fastapi import FastAPI, WebSocket
from fastapi.testclient import TestClient

app = FastAPI()


@app.websocket("/ws/types")
async def types(websocket: WebSocket) -> None:
    await websocket.accept()
    text = await websocket.receive_text()
    await websocket.send_text(text.upper())
    blob = await websocket.receive_bytes()
    await websocket.send_bytes(blob[::-1])
    payload = await websocket.receive_json()
    await websocket.send_json({"title": payload["title"], "length": len(payload["title"])})
    # JSON over binary frames
    payload = await websocket.receive_json(mode="binary")
    await websocket.send_json({"binary": True, **payload}, mode="binary")
    # the raw ASGI message, whichever kind it is
    message = await websocket.receive()
    await websocket.send_json({"keys": sorted(k for k, v in message.items() if v is not None)})
    await websocket.close()


client = TestClient(app)
with client.websocket_connect("/ws/types") as ws:
    ws.send_text("dune")
    print(ws.receive_text())
    ws.send_bytes(b"\x01\x02\x03")
    print(ws.receive_bytes())
    ws.send_json({"title": "Dune"})
    print(ws.receive_json())
    ws.send_json({"a": 1}, mode="binary")
    print(ws.receive_json(mode="binary"))
    ws.send_text("raw")
    print(ws.receive_json())

DUNE
b'\x03\x02\x01'
{'title': 'Dune', 'length': 4}
{'binary': True, 'a': 1}
{'keys': ['text', 'type']}

from typing import Annotated

from fastapi import (Cookie, Depends, FastAPI, Header, Query, WebSocket,
                     WebSocketDisconnect, WebSocketException, status)
from fastapi.testclient import TestClient

app = FastAPI()


async def get_username(
    websocket: WebSocket,
    session: Annotated[str | None, Cookie()] = None,
    token: Annotated[str | None, Query()] = None,
) -> str:
    if session is None and token is None:
        raise WebSocketException(code=status.WS_1008_POLICY_VIOLATION, reason="credentials required")
    return f"reader-{(session or token)[:3]}"


Username = Annotated[str, Depends(get_username)]


@app.websocket("/ws/books/{book_id}")
async def book_socket(
    websocket: WebSocket,
    book_id: int,
    username: Username,
    last_seen: Annotated[int, Query(ge=0)] = 0,
    user_agent: Annotated[str | None, Header()] = None,
) -> None:
    await websocket.accept()
    await websocket.send_json(
        {"book_id": book_id, "who": username, "last_seen": last_seen, "agent": user_agent}
    )
    await websocket.close()


client = TestClient(app)


def attempt(url: str, **kw) -> None:
    try:
        with client.websocket_connect(url, **kw) as ws:
            print("OK  ", url, ws.receive_json())
            try:
                ws.receive_text()
            except WebSocketDisconnect as exc:
                print("     then closed", exc.code)
    except WebSocketDisconnect as exc:
        print("DENY", url, "code", exc.code, "reason", repr(exc.reason))


attempt("/ws/books/7?token=abc123&last_seen=4", headers={"user-agent": "demo-client"})
attempt("/ws/books/7")
attempt("/ws/books/not-a-number?token=abc123")
attempt("/ws/books/7?token=abc123&last_seen=-1")
client.cookies.set("session", "xyz789")
attempt("/ws/books/9")

OK   /ws/books/7?token=abc123&last_seen=4 {'book_id': 7, 'who': 'reader-abc', 'last_seen': 4, 'agent': 'demo-client'}
     then closed 1000
DENY /ws/books/7 code 1008 reason 'credentials required'
DENY /ws/books/not-a-number?token=abc123 code 1008 reason [{'type': 'int_parsing', 'loc': ['path', 'book_id'], 'msg': 'Input should be a valid integer, unable to parse string as an integer', 'input': 'not-a-number'}]
DENY /ws/books/7?token=abc123&last_seen=-1 code 1008 reason [{'type': 'greater_than_equal', 'loc': ['query', 'last_seen'], 'msg': 'Input should be greater than or equal to 0', 'input': '-1', 'ctx': {'ge': 0}}]
OK   /ws/books/9 {'book_id': 9, 'who': 'reader-xyz', 'last_seen': 0, 'agent': 'testclient'}
     then closed 1000

from fastapi import FastAPI, WebSocket, WebSocketDisconnect, status
from fastapi.exceptions import WebSocketRequestValidationError
from fastapi.testclient import TestClient

app = FastAPI()


@app.exception_handler(WebSocketRequestValidationError)
async def ws_validation_handler(websocket: WebSocket, exc: WebSocketRequestValidationError):
    await websocket.close(code=status.WS_1008_POLICY_VIOLATION, reason="invalid parameters")


@app.websocket("/ws/books/{book_id}")
async def book_socket(websocket: WebSocket, book_id: int) -> None:
    await websocket.accept()
    await websocket.send_json({"book_id": book_id})
    await websocket.close()


client = TestClient(app)
try:
    with client.websocket_connect("/ws/books/abc"):
        pass
except WebSocketDisconnect as exc:
    print(exc.code, repr(exc.reason))

1008 'invalid parameters'

from fastapi import FastAPI, WebSocket, WebSocketException, status

app = FastAPI()


@app.websocket("/ws/deny-before-accept")
async def deny(websocket: WebSocket):
    await websocket.close(code=status.WS_1008_POLICY_VIOLATION)


@app.websocket("/ws/exception-after-accept")
async def exc_after(websocket: WebSocket):
    await websocket.accept()
    raise WebSocketException(code=status.WS_1008_POLICY_VIOLATION, reason="nope")


@app.websocket("/ws/raw-loop")
async def raw_loop(websocket: WebSocket):
    await websocket.accept()
    while True:
        text = await websocket.receive_text()
        await websocket.send_text(text)


@app.websocket("/ws/books/{book_id}")
async def typed(websocket: WebSocket, book_id: int):
    await websocket.accept()
    await websocket.close()

@app.websocket("/ws/boom")
async def boom(websocket: WebSocket):
    await websocket.accept()
    raise RuntimeError("bug")

import asyncio, websockets

import os
BASE = "ws://127.0.0.1:" + os.environ.get("PORT", "18801")


async def main():
    for path in ["/ws/deny-before-accept", "/ws/books/abc", "/ws/nope"]:
        try:
            async with websockets.connect(BASE + path):
                print(path, "connected")
        except Exception as e:
            print(path, type(e).__name__, e)
    async with websockets.connect(BASE + "/ws/exception-after-accept") as ws:
        try:
            await ws.recv()
        except websockets.ConnectionClosed as e:
            print("after-accept exception:", ws.close_code, repr(ws.close_reason))
    async with websockets.connect(BASE + "/ws/boom") as ws:
        try:
            await ws.recv()
        except websockets.ConnectionClosed as e:
            print("boom:", ws.close_code, repr(ws.close_reason))
    async with websockets.connect(BASE + "/ws/raw-loop") as ws:
        await ws.send("hi"); print(await ws.recv())
    print("raw-loop client left")

asyncio.run(main())

/ws/deny-before-accept InvalidStatus server rejected WebSocket connection: HTTP 403
/ws/books/abc InvalidStatus server rejected WebSocket connection: HTTP 403
/ws/nope InvalidStatus server rejected WebSocket connection: HTTP 403
after-accept exception: 1008 'nope'
boom: 1006 ''
hi
raw-loop client left

INFO:     127.0.0.1:60055 - "WebSocket /ws/deny-before-accept" 403
INFO:     connection rejected (403 Forbidden)
INFO:     127.0.0.1:60059 - "WebSocket /ws/boom" [accepted]
ERROR:    Exception in ASGI application
RuntimeError: bug
INFO:     127.0.0.1:60060 - "WebSocket /ws/raw-loop" [accepted]
ERROR:    Exception in ASGI application
starlette.websockets.WebSocketDisconnect: (1000, '')

import asyncio
from collections import defaultdict
from datetime import datetime, timezone
from typing import Annotated, Literal

from fastapi import FastAPI, Query, WebSocket
from pydantic import BaseModel, Field, ValidationError


class ChatIn(BaseModel):
    text: Annotated[str, Field(min_length=1, max_length=500)]


class ChatEvent(BaseModel):
    kind: Literal["join", "message", "leave", "error"]
    user: str
    text: str = ""
    at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))


class ConnectionManager:
    """Tracks the open sockets of each book's chat room, inside one process."""

    def __init__(self) -> None:
        self._rooms: dict[int, set[WebSocket]] = defaultdict(set)

    async def connect(self, book_id: int, websocket: WebSocket) -> None:
        await websocket.accept()
        self._rooms[book_id].add(websocket)

    def disconnect(self, book_id: int, websocket: WebSocket) -> None:
        room = self._rooms.get(book_id)
        if room is not None:
            room.discard(websocket)
            if not room:
                del self._rooms[book_id]

    def size(self, book_id: int) -> int:
        return len(self._rooms.get(book_id, ()))

    async def _send(self, book_id: int, websocket: WebSocket, data: str) -> None:
        try:
            await asyncio.wait_for(websocket.send_text(data), timeout=2.0)
        except Exception:  # dead or too slow: drop it, keep the room alive
            self.disconnect(book_id, websocket)

    async def broadcast(self, book_id: int, event: ChatEvent) -> None:
        await self.broadcast_raw(book_id, event.model_dump_json())

    async def broadcast_raw(self, book_id: int, data: str) -> None:
        sockets = list(self._rooms.get(book_id, ()))  # snapshot: the set changes while we await
        await asyncio.gather(*(self._send(book_id, ws, data) for ws in sockets))


manager = ConnectionManager()
app = FastAPI()


@app.websocket("/ws/books/{book_id}/chat")
async def book_chat(
    websocket: WebSocket, book_id: int, user: Annotated[str, Query(min_length=1, max_length=30)]
) -> None:
    await manager.connect(book_id, websocket)
    await manager.broadcast(book_id, ChatEvent(kind="join", user=user))
    try:
        async for raw in websocket.iter_json():
            try:
                incoming = ChatIn.model_validate(raw)
            except ValidationError:
                await websocket.send_text(
                    ChatEvent(kind="error", user="server", text="text must be 1-500 characters").model_dump_json()
                )
                continue
            await manager.broadcast(book_id, ChatEvent(kind="message", user=user, text=incoming.text))
    finally:
        manager.disconnect(book_id, websocket)
        await manager.broadcast(book_id, ChatEvent(kind="leave", user=user))

import json

from fastapi.testclient import TestClient

from chat import app, manager


def show(label, ws):
    event = json.loads(ws.receive_text())
    print(f"{label} <- {event['kind']:8} {event['user']:6} {event['text']!r}")


with TestClient(app) as client:  # one shared event loop for every connection
    with client.websocket_connect("/ws/books/7/chat?user=ana") as ana:
        show("ana", ana)
        with client.websocket_connect("/ws/books/7/chat?user=ben") as ben:
            show("ana", ana)
            show("ben", ben)
            with client.websocket_connect("/ws/books/8/chat?user=eve") as eve:
                show("eve", eve)  # a different book: sees only her own join
                ben.send_json({"text": "Loved the ending"})
                show("ana", ana)
                show("ben", ben)
                ana.send_json({"text": ""})
                show("ana", ana)
                print("room 7:", manager.size(7), "room 8:", manager.size(8))
print("after:", manager.size(7), manager.size(8))

ana <- join     ana    ''
ana <- join     ben    ''
ben <- join     ben    ''
eve <- join     eve    ''
ana <- message  ben    'Loved the ending'
ben <- message  ben    'Loved the ending'
ana <- error    server 'text must be 1-500 characters'
room 7: 2 room 8: 1
after: 0 0

import asyncio, json, os, websockets

BASE = "ws://127.0.0.1:" + os.environ["PORT"]


async def show(label, ws):
    d = json.loads(await ws.recv())
    print(f"{label} <- {d['kind']:8} {d['user']:4} {d['text']!r}")


async def main():
    async with websockets.connect(f"{BASE}/ws/books/7/chat?user=ana") as ana:
        await show("ana", ana)
        async with websockets.connect(f"{BASE}/ws/books/7/chat?user=ben") as ben:
            await show("ana", ana); await show("ben", ben)
            await ben.send(json.dumps({"text": "Loved the ending"}))
            await show("ana", ana); await show("ben", ben)
        await show("ana", ana)  # ben left


asyncio.run(main())

ana <- join     ana  ''
ana <- join     ben  ''
ben <- join     ben  ''
ana <- message  ben  'Loved the ending'
ben <- message  ben  'Loved the ending'
ana <- leave    ben  ''

import time
from collections import deque

import anyio
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, status
from fastapi.testclient import TestClient

IDLE_SECONDS = 0.3        # 60 or more in production
MAX_MESSAGES = 3          # per window
WINDOW_SECONDS = 1.0
MAX_TEXT_BYTES = 1024


class SlidingWindow:
    """At most `limit` events per `window` seconds, for one connection."""

    def __init__(self, limit: int, window: float) -> None:
        self.limit, self.window = limit, window
        self.events: deque[float] = deque()

    def allow(self) -> bool:
        now = time.monotonic()
        while self.events and now - self.events[0] > self.window:
            self.events.popleft()
        if len(self.events) >= self.limit:
            return False
        self.events.append(now)
        return True


app = FastAPI()


@app.websocket("/ws/guarded")
async def guarded(websocket: WebSocket) -> None:
    await websocket.accept()
    limiter = SlidingWindow(MAX_MESSAGES, WINDOW_SECONDS)
    while True:
        try:
            with anyio.fail_after(IDLE_SECONDS):
                text = await websocket.receive_text()
        except TimeoutError:
            await websocket.close(code=status.WS_1001_GOING_AWAY, reason="idle timeout")
            return
        except WebSocketDisconnect:
            return
        if len(text.encode()) > MAX_TEXT_BYTES:
            await websocket.close(code=status.WS_1009_MESSAGE_TOO_BIG, reason="message too big")
            return
        if not limiter.allow():
            await websocket.close(code=status.WS_1008_POLICY_VIOLATION, reason="rate limit exceeded")
            return
        await websocket.send_text(f"ok: {text}")


client = TestClient(app)


def run(label, messages, wait=0.0):
    try:
        with client.websocket_connect("/ws/guarded") as ws:
            for m in messages:
                ws.send_text(m)
                print(label, ws.receive_text())
            time.sleep(wait)
            ws.receive_text()
    except WebSocketDisconnect as exc:
        print(label, "closed", exc.code, exc.reason)


run("flood", ["a", "b", "c", "d"])
run("big  ", ["x" * 2000])
run("idle ", ["hello"], wait=0.6)

flood ok: a
flood ok: b
flood ok: c
flood closed 1008 rate limit exceeded
big   closed 1009 message too big
idle  ok: hello
idle  closed 1001 idle timeout

import asyncio
import contextlib
from collections.abc import AsyncIterator
from typing import Annotated

import redis.asyncio as redis
from fastapi import FastAPI, Query, WebSocket

from chat import ChatEvent, ChatIn, ConnectionManager  # the per-process manager from the previous section


class ChatHub:
    """Publishes every chat event to Redis; each process relays what it hears to its own sockets."""

    def __init__(self, client: redis.Redis, local: ConnectionManager) -> None:
        self.client = client
        self.local = local
        self._pubsub = client.pubsub()
        self._task: asyncio.Task[None] | None = None

    async def start(self) -> None:
        await self._pubsub.psubscribe("book:*")
        self._task = asyncio.create_task(self._relay())

    async def stop(self) -> None:
        if self._task is not None:
            self._task.cancel()
            with contextlib.suppress(asyncio.CancelledError):
                await self._task
        await self._pubsub.aclose()

    async def publish(self, book_id: int, event: ChatEvent) -> None:
        await self.client.publish(f"book:{book_id}", event.model_dump_json())

    async def _relay(self) -> None:
        async for message in self._pubsub.listen():
            if message["type"] != "pmessage":
                continue
            book_id = int(message["channel"].removeprefix("book:"))
            await self.local.broadcast_raw(book_id, message["data"])


def create_app(client: redis.Redis) -> FastAPI:
    @contextlib.asynccontextmanager
    async def lifespan(app: FastAPI) -> AsyncIterator[dict]:
        hub = ChatHub(client, ConnectionManager())
        await hub.start()
        yield {"hub": hub}
        await hub.stop()

    app = FastAPI(lifespan=lifespan)

    @app.websocket("/ws/books/{book_id}/chat")
    async def chat(websocket: WebSocket, book_id: int,
                   user: Annotated[str, Query(min_length=1, max_length=30)]) -> None:
        hub: ChatHub = websocket.state.hub
        await hub.local.connect(book_id, websocket)
        await hub.publish(book_id, ChatEvent(kind="join", user=user))
        try:
            async for raw in websocket.iter_json():
                incoming = ChatIn.model_validate(raw)
                await hub.publish(book_id, ChatEvent(kind="message", user=user, text=incoming.text))
        finally:
            hub.local.disconnect(book_id, websocket)
            await hub.publish(book_id, ChatEvent(kind="leave", user=user))

    return app

import json
import fakeredis
from fakeredis import FakeAsyncRedis
from fastapi.testclient import TestClient
from hub_chat import create_app

server = fakeredis.FakeServer()  # stands in for one shared Redis; two apps = two workers
worker_a = create_app(FakeAsyncRedis(server=server, decode_responses=True))
worker_b = create_app(FakeAsyncRedis(server=server, decode_responses=True))


def show(label, ws):
    d = json.loads(ws.receive_text())
    print(f"{label} <- {d['kind']:8} {d['user']:4} {d['text']!r}")


with TestClient(worker_a) as a, TestClient(worker_b) as b:
    with a.websocket_connect("/ws/books/7/chat?user=ana") as ana:
        show("ana@A", ana)
        with b.websocket_connect("/ws/books/7/chat?user=ben") as ben:
            show("ana@A", ana); show("ben@B", ben)
            ben.send_json({"text": "Loved the ending"})
            show("ana@A", ana); show("ben@B", ben)
            ana.send_json({"text": "Same here"})
            show("ana@A", ana); show("ben@B", ben)

ana@A <- join     ana  ''
ana@A <- join     ben  ''
ben@B <- join     ben  ''
ana@A <- message  ben  'Loved the ending'
ben@B <- message  ben  'Loved the ending'
ana@A <- message  ana  'Same here'
ben@B <- message  ana  'Same here'

location /ws/ {
    proxy_pass http://backend;
    proxy_set_header Upgrade $http_upgrade;
    proxy_set_header Connection "upgrade";
    proxy_read_timeout 120s;
}

import secrets

import pytest
from fastapi.testclient import TestClient
from starlette.websockets import WebSocketDisconnect

from auth_chat import Settings, app, get_settings
from chat import app as room_app

TEST_SETTINGS = Settings(chat_ticket_secret=secrets.token_hex(32), chat_ticket_ttl_seconds=30)


@pytest.fixture
def client():
    app.dependency_overrides[get_settings] = lambda: TEST_SETTINGS
    with TestClient(app) as c:
        yield c
    app.dependency_overrides.clear()


def ticket_for(client, book_id, user="ana"):
    response = client.post(f"/books/{book_id}/chat-ticket", headers={"Authorization": f"Bearer user:{user}"})
    assert response.status_code == 200
    return response.json()["ticket"]


def test_connect_send_receive(client):
    with client.websocket_connect(f"/ws/books/7/chat?ticket={ticket_for(client, 7)}") as ws:
        assert ws.receive_json() == {"welcome": "ana", "book": 7}
        ws.send_json({"text": "hi"})
        assert ws.receive_json() == {"user": "ana", "echo": {"text": "hi"}}


def test_ticket_for_another_book_is_refused(client):
    with pytest.raises(WebSocketDisconnect) as exc:
        with client.websocket_connect(f"/ws/books/8/chat?ticket={ticket_for(client, 7)}"):
            pass
    assert exc.value.code == 1008
    assert exc.value.reason == "ticket is for another book"


def test_foreign_origin_is_refused(client):
    with pytest.raises(WebSocketDisconnect) as exc:
        with client.websocket_connect(f"/ws/books/7/chat?ticket={ticket_for(client, 7)}",
                                      headers={"origin": "https://evil.example.net"}):
            pass
    assert exc.value.code == 1008


def test_room_broadcast_reaches_everyone_in_the_book():
    with TestClient(room_app) as c:  # one shared event loop, so two sockets can talk
        with c.websocket_connect("/ws/books/1/chat?user=ana") as ana, \
             c.websocket_connect("/ws/books/1/chat?user=ben") as ben:
            ana.receive_json(); ana.receive_json(); ben.receive_json()  # join events
            ben.send_json({"text": "hello"})
            assert ana.receive_json()["text"] == "hello"
            assert ben.receive_json()["text"] == "hello"

....                                                                     [100%]
4 passed in 0.41s