FastAPI WebSocket 本番実装ガイド：双方向リアルタイム通信を接続管理・認証・水平スケールまで作り込む

from fastapi import FastAPI, WebSocket

app = FastAPI()

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    # ハンドシェイクを受理する。これを呼ぶまで通信は始まらない。
    await websocket.accept()
    while True:
        # クライアントからの1メッセージを待つ（届くまでここで非同期に待機）
        data = await websocket.receive_text()
        # そのまま送り返す（エコー）
        await websocket.send_text(f"Message text was: {data}")

from fastapi import FastAPI, WebSocket, WebSocketDisconnect

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"echo: {data}")
    except WebSocketDisconnect:
        # クライアントが切断した。ここに来たら接続はもう無い。
        # ここでクリーンアップ（後述の ConnectionManager からの除去など）を行う。
        ...

from typing import Literal
from pydantic import BaseModel, ValidationError, Field

# 受信メッセージの「許される形」を型で定義する（判別可能なユニオン）
class ChatMessage(BaseModel):
    type: Literal["chat"]
    text: str = Field(min_length=1, max_length=2000)   # 長さ上限＝防御

class PingMessage(BaseModel):
    type: Literal["ping"]

# type で分岐できるよう discriminated union にする
IncomingMessage = ChatMessage | PingMessage

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    try:
        while True:
            raw = await websocket.receive_json()    # まずは dict として受ける
            try:
                # ここが境界検証。未知の type・型違い・過大長はここで弾かれる。
                message = _parse_incoming(raw)
            except ValidationError:
                # 不正な入力で接続ごと殺さず、エラーを返して接続は維持する判断もある
                await websocket.send_json({"type": "error", "detail": "invalid message"})
                continue
            await _handle(websocket, message)
    except WebSocketDisconnect:
        ...

from pydantic import TypeAdapter

# TypeAdapter はユニオン型をそのまま検証できる（モデルでなくてもよい）
_incoming_adapter: TypeAdapter[IncomingMessage] = TypeAdapter(IncomingMessage)

def _parse_incoming(raw: object) -> IncomingMessage:
    # 外部入力 → 型の確定した内部表現へ。以降のコードは型を信頼できる。
    return _incoming_adapter.validate_python(raw)

from fastapi import WebSocket

class ConnectionManager:
    def __init__(self) -> None:
        # 現在アクティブな接続の集合。プロセスのメモリ上に持つ（←この前提が第6章の論点）。
        self.active_connections: list[WebSocket] = []

    async def connect(self, websocket: WebSocket) -> None:
        await websocket.accept()                 # 受理してから
        self.active_connections.append(websocket)  # 台帳に載せる

    def disconnect(self, websocket: WebSocket) -> None:
        self.active_connections.remove(websocket)  # 台帳から外す

    async def send_personal_message(self, message: str, websocket: WebSocket) -> None:
        await websocket.send_text(message)         # 特定の1接続へ

    async def broadcast(self, message: str) -> None:
        for connection in self.active_connections:
            await connection.send_text(message)    # 全接続へ

manager = ConnectionManager()

from fastapi import WebSocket, WebSocketDisconnect

@app.websocket("/ws/{client_id}")
async def websocket_endpoint(websocket: WebSocket, client_id: int):
    await manager.connect(websocket)
    try:
        while True:
            data = await websocket.receive_text()
            # 自分以外も含む全員へ配信（部屋・チャンネル単位にするなら台帳を分ける）
            await manager.broadcast(f"Client #{client_id} says: {data}")
    except WebSocketDisconnect:
        manager.disconnect(websocket)
        await manager.broadcast(f"Client #{client_id} left the chat")

import asyncio
import logging

logger = logging.getLogger("ws")

class ConnectionManager:
    def __init__(self) -> None:
        self.active_connections: set[WebSocket] = set()   # 重複登録を避けるなら set

    async def connect(self, websocket: WebSocket) -> None:
        await websocket.accept()
        self.active_connections.add(websocket)

    def disconnect(self, websocket: WebSocket) -> None:
        # discard は「居なくても例外を投げない」。二重切断でも安全（冪等）。
        self.active_connections.discard(websocket)

    async def broadcast(self, payload: dict) -> None:
        # 送信中に1つでも死んでいる接続があると例外で全体が止まる。
        # gather(return_exceptions=True) で個別の失敗を隔離する。
        dead: list[WebSocket] = []
        results = await asyncio.gather(
            *(conn.send_json(payload) for conn in self.active_connections),
            return_exceptions=True,
        )
        for conn, result in zip(self.active_connections, results):
            if isinstance(result, Exception):
                dead.append(conn)              # 送れなかった接続は死んでいる
        for conn in dead:
            self.disconnect(conn)              # まとめて掃除（クリーンアップ）

@app.websocket("/ws/{client_id}")
async def websocket_endpoint(websocket: WebSocket, client_id: int):
    await manager.connect(websocket)
    try:
        while True:
            raw = await websocket.receive_json()
            message = _parse_incoming(raw)        # 第3章の境界検証
            await manager.broadcast({"from": client_id, "text": message.text})
    except WebSocketDisconnect:
        pass                                       # 正常な切断
    finally:
        # どんな経路で抜けても必ず台帳から外す（接続リーク防止）
        manager.disconnect(websocket)

from typing import Annotated
from fastapi import (
    FastAPI, WebSocket, WebSocketException, WebSocketDisconnect,
    Cookie, Query, Depends, status,
)

app = FastAPI()

async def get_token(
    websocket: WebSocket,
    # ブラウザの制約上、ヘッダは使いにくい → Cookie かクエリで受ける
    session: Annotated[str | None, Cookie()] = None,
    token: Annotated[str | None, Query()] = None,
) -> str:
    candidate = session or token
    if candidate is None:
        # 認証情報が無い → ポリシー違反として接続を拒否（accept 自体が行われない）
        raise WebSocketException(code=status.WS_1008_POLICY_VIOLATION)
    return candidate

@app.websocket("/ws")
async def websocket_endpoint(
    websocket: WebSocket,
    token: Annotated[str, Depends(get_token)],   # ここで認証が走る
):
    user = _verify_jwt(token)                    # JWT 検証（次項）。失敗なら例外で弾く
    await websocket.accept()                     # 認証を通って初めて受理
    try:
        while True:
            data = await websocket.receive_text()
            await websocket.send_text(f"hello {user.username}: {data}")
    except WebSocketDisconnect:
        ...

ALLOWED_ORIGINS = {"https://app.example.com"}   # 環境変数から読む

async def check_origin(websocket: WebSocket) -> None:
    origin = websocket.headers.get("origin")
    # 許可されていないオリジンからの接続は拒否
    if origin not in ALLOWED_ORIGINS:
        raise WebSocketException(code=status.WS_1008_POLICY_VIOLATION)

[クライアントA] ── レプリカ#1 ┐                      ┌ レプリカ#1 → [クライアントA]
                              ├─→ Redis (publish) ──┤
[クライアントB] ── レプリカ#2 ┘    channel: "room:42" └ レプリカ#2 → [クライアントB]

import asyncio
import json
from redis.asyncio import Redis

class BroadcastHub:
    """ローカル接続の管理 + Redis Pub/Sub での全レプリカ配信。"""

    def __init__(self, redis: Redis, channel: str = "broadcast") -> None:
        self._redis = redis
        self._channel = channel
        self._local: set[WebSocket] = set()     # このレプリカが持つ接続だけ

    async def connect(self, websocket: WebSocket) -> None:
        await websocket.accept()
        self._local.add(websocket)

    def disconnect(self, websocket: WebSocket) -> None:
        self._local.discard(websocket)

    async def publish(self, payload: dict) -> None:
        # 直接ローカルへ送らず、まず Redis に publish（全レプリカへ波及させる）
        await self._redis.publish(self._channel, json.dumps(payload))

    async def _fan_out_to_local(self, payload: dict) -> None:
        # Redis から受けたメッセージを、このレプリカのローカル接続へ配る
        dead = []
        for conn in self._local:
            try:
                await conn.send_json(payload)
            except Exception:
                dead.append(conn)
        for conn in dead:
            self.disconnect(conn)

    async def run_subscriber(self) -> None:
        # 起動時に1つだけ走らせる常駐タスク：Redis を購読し続ける
        pubsub = self._redis.pubsub()
        await pubsub.subscribe(self._channel)
        async for message in pubsub.listen():
            if message["type"] == "message":
                await self._fan_out_to_local(json.loads(message["data"]))

from contextlib import asynccontextmanager

@asynccontextmanager
async def lifespan(app: FastAPI):
    redis = Redis.from_url(settings.redis_url)   # URL は環境変数
    hub = BroadcastHub(redis)
    task = asyncio.create_task(hub.run_subscriber())   # 購読を常駐させる
    app.state.hub = hub
    try:
        yield
    finally:
        task.cancel()                            # グレースフルに停止
        await redis.aclose()

import time
from collections import deque

class RateLimiter:
    """1接続あたり: 直近 window 秒で max_messages を超えたら True（＝超過）。"""
    def __init__(self, max_messages: int = 20, window: float = 1.0) -> None:
        self._max = max_messages
        self._window = window
        self._timestamps: deque[float] = deque()

    def is_exceeded(self) -> bool:
        now = time.monotonic()
        while self._timestamps and now - self._timestamps[0] > self._window:
            self._timestamps.popleft()         # 古い記録を捨てる（スライディングウィンドウ）
        self._timestamps.append(now)
        return len(self._timestamps) > self._max

limiter = RateLimiter()
# ループ内で:
if limiter.is_exceeded():
    await websocket.close(code=1008, reason="rate limit exceeded")
    break

import asyncio

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    try:
        while True:
            try:
                # 一定時間メッセージが無ければタイムアウト → 生存確認・切断判断へ
                data = await asyncio.wait_for(websocket.receive_text(), timeout=30.0)
            except asyncio.TimeoutError:
                await websocket.send_json({"type": "ping"})   # 生きてるか問い合わせ
                continue
            await _handle(websocket, data)
    except WebSocketDisconnect:
        ...
    finally:
        manager.disconnect(websocket)

# 概念コード（クライアント側のロジック例）。実装言語は問わない。
def next_delay(attempt: int) -> float:
    base = min(2 ** attempt, 30)           # 1, 2, 4, ... 最大30秒で頭打ち
    return base * (0.5 + random.random())  # ジッタで殺到を散らす

from fastapi.testclient import TestClient
from app.main import app

client = TestClient(app)

def test_echo():
    # with でハンドシェイク〜切断まで面倒を見てくれる
    with client.websocket_connect("/ws") as websocket:
        websocket.send_text("hello")
        data = websocket.receive_text()
        assert data == "echo: hello"

def test_json_roundtrip():
    with client.websocket_connect("/ws") as websocket:
        websocket.send_json({"type": "chat", "text": "hi"})
        assert websocket.receive_json() == {"from": 1, "text": "hi"}

import pytest
from fastapi import status
from starlette.websockets import WebSocketDisconnect

def test_invalid_message_does_not_crash():
    with client.websocket_connect("/ws") as websocket:
        websocket.send_json({"type": "unknown"})         # 未知の type
        res = websocket.receive_json()
        assert res["type"] == "error"                    # 接続は維持しつつエラーを返す

def test_missing_token_is_rejected():
    # 認証無しでの接続は、ハンドシェイク段階で拒否される
    with pytest.raises(WebSocketDisconnect) as exc:
        with client.websocket_connect("/ws") as websocket:
            websocket.receive_text()
    assert exc.value.code == status.WS_1008_POLICY_VIOLATION