WebSocket

前言

WebSocket 是全双工通信协议：连接建立后，服务端和客户端可以随时互相推送消息，适合聊天、实时通知、协同编辑等场景。FastAPI 原生支持 WebSocket。

本文示例已验证于 FastAPI 0.141.1。

基础用法

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:
        print("客户端断开连接")

常用方法：

• await websocket.accept()：接受连接。
• await websocket.receive_text() / receive_bytes() / receive_json()：接收文本/二进制/JSON。
• await websocket.send_text(...) / send_bytes(...) / send_json(...)：发送文本/二进制/JSON。
• await websocket.close()：关闭连接。
• WebSocketDisconnect：客户端断开时抛出，用于清理。

连接管理 + 广播

实际应用中通常需要维护所有在线连接，实现群发（广播）：

from fastapi import FastAPI, WebSocket, WebSocketDisconnect

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):
        if websocket in self.active_connections:
            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(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("有人离开了聊天室")

WebSocket 中使用依赖注入

WebSocket 路由同样支持 Depends（如认证依赖）：

from fastapi import Cookie, Depends, WebSocket


async def get_current_user(token: str = Cookie(...)) -> str:
    # 实际项目应校验 token（JWT 等），这里仅做演示
    return token


@app.websocket("/ws/auth")
async def auth_ws(websocket: WebSocket, user: str = Depends(get_current_user)):
    await websocket.accept()
    await websocket.send_text(f"欢迎, {user}")

客户端测试

使用 Starlette 的 TestClient 可以方便地测试 WebSocket 接口：

from fastapi.testclient import TestClient

client = TestClient(app)

with client.websocket_connect("/ws") as websocket:
    websocket.send_text("hello")
    data = websocket.receive_text()
    print(data)  # echo: hello

注意事项

• 生产环境建议在 Nginx / Caddy 等反向代理上配置 WebSocket 升级头（Upgrade / Connection），见 服务部署。
• 广播时单个慢客户端会拖慢整体发送，可引入队列/并发控制；多进程部署时广播需要借助 Redis Pub/Sub 等外部机制（进程间内存不共享）。
• 心跳保活：长时间无消息的连接可能被中间设备断开，服务端可定时发送 ping 或与客户端约定心跳。