FastAPI WebSocket 基础指南

什么是 WebSocket

WebSocket 是一种在单个 TCP 连接上进行全双工通信的协议。与 HTTP 的「请求-响应」模式不同，WebSocket 连接一旦建立，客户端和服务端可以随时互相发送消息，非常适合聊天应用、实时通知、股票行情、协作编辑等场景。

FastAPI 对 WebSocket 提供了原生支持，使用起来非常简洁。

最简单的 WebSocket 端点

from fastapi import FastAPI, WebSocket

app = FastAPI()


@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    while True:
        data = await websocket.receive_text()
        await websocket.send_text(f"收到消息: {data}")

要点说明：

• 用 @app.websocket() 装饰器定义 WebSocket 路由
• 必须先调用 await websocket.accept() 接受连接，否则无法通信
• receive_text() 接收文本消息，send_text() 发送文本消息

WebSocket 对象的核心方法

处理异常断开

客户端可能随时断开连接，捕获异常是必要的：

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"回复: {data}")
    except WebSocketDisconnect:
        print("客户端已断开连接")

FastAPI 内置了 WebSocketDisconnect 异常，当客户端断开或网络异常时会自动抛出，无需手动判断。

连接管理：携带查询参数

WebSocket 连接建立时可以携带查询参数，常用于传递 token 或房间号：

@app.websocket("/ws/{client_id}")
async def websocket_endpoint(websocket: WebSocket, client_id: str, token: str = ""):
    # 可以在此处验证 token
    await websocket.accept()
    try:
        while True:
            data = await websocket.receive_text()
            await websocket.send_text(f"[{client_id}] 说: {data}")
    except WebSocketDisconnect:
        print(f"{client_id} 已离开")

访问时使用 ws://localhost:8000/ws/alice?token=abc123，FastAPI 会自动解析路径参数和查询参数。

广播消息：多客户端聊天室

实际场景中常常需要将消息推送给所有连接的客户端，需要维护一个连接池：

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from typing import List

app = FastAPI()


class ConnectionManager:
    def __init__(self):
        self.active_connections: List[WebSocket] = []

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)

    def disconnect(self, websocket: WebSocket):
        self.active_connections.remove(websocket)

    async def broadcast(self, message: str):
        for connection in self.active_connections:
            await connection.send_text(message)


manager = ConnectionManager()


@app.websocket("/ws/chat")
async def chat_endpoint(websocket: WebSocket):
    await manager.connect(websocket)
    try:
        while True:
            data = await websocket.receive_text()
            await manager.broadcast(f"有人发来消息: {data}")
    except WebSocketDisconnect:
        manager.disconnect(websocket)
        await manager.broadcast("一位用户离开了聊天室")

这个 ConnectionManager 封装了连接管理逻辑：连接时注册、断开时移除、广播时遍历所有连接发送消息。

在生产环境中，如果有多个 worker 进程，消息广播需要借助 Redis Pub/Sub 或消息队列来实现跨进程通信。

依赖注入与 WebSocket

WebSocket 端点同样可以使用 Depends 进行依赖注入，比如验证 token：

from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Depends, Query


async def get_token(token: str = Query(...)):
    # 此处可以查询数据库验证 token 有效性
    if token != "secret":
        raise Exception("认证失败")
    return token


@app.websocket("/ws/secure")
async def secure_websocket(
    websocket: WebSocket,
    token: str = Depends(get_token),
):
    await websocket.accept()
    try:
        while True:
            data = await websocket.receive_text()
            await websocket.send_text(f"已认证用户: {data}")
    except WebSocketDisconnect:
        print("断开连接")

注意：依赖在 accept() 之前执行，如果认证失败抛出异常，连接不会被接受，客户端会收到 HTTP 403 响应。

生产环境部署注意事项

部署 WebSocket 应用时需要注意以下几点：

1. 反向代理配置：Nginx 需要升级 HTTP 连接为 WebSocket，添加 proxy_set_header Upgrade $http_upgrade 和 proxy_set_header Connection "upgrade"
2. 超时设置：WebSocket 是长连接，需要调整代理和负载均衡器的超时时间
3. 多进程问题：使用 ASGI 服务器（如 Uvicorn）时，多 worker 下连接分布在不同进程，广播需借助 Redis 等外部工具
4. 心跳检测：长时间无消息时连接可能被中间件断开，建议定期发送 ping/pong 保活

小结

FastAPI 的 WebSocket 支持非常直观：用装饰器定义端点，用 accept() 建立连接，用 receive_* 和 send_* 收发消息。配合依赖注入和连接管理器，可以轻松构建实时通信功能。核心要记住：先 accept，再收发，最后处理 WebSocketDisconnect 异常即可。