FastAPI WebSocket实时通信完全指南

📂 所属阶段：第六阶段 — 2026 特色专题（AI 集成篇）
🔗 相关章节：流式响应 StreamingResponse · OAuth2 与 JWT 鉴权

目录

WebSocket基础概念

为什么用WebSocket？

传统HTTP是短连接单向轮询，每次请求都要带完整Header，浪费带宽延迟高。WebSocket是长连接双向通信，仅需一次握手，开销极低适合实时场景。

核心应用场景

  • AI实时流式对话、多人聊天室
  • 实时通知、协作编辑状态同步
  • 游戏、IoT设备监控

FastAPI WebSocket入门

基础单连接端点

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse
import json
from datetime import datetime

app = FastAPI(title="道满WebSocket入门")

# 简单测试页面
html = """
<!DOCTYPE html>
<html>
<head>
    <title>入门测试</title>
    <style>body{font-family:Arial;margin:20px}#msg{height:300px;overflow-y:auto;border:1px solid #ccc;padding:10px}</style>
</head>
<body>
    <h1>入门测试</h1>
    <div id="status">未连接</div>
    <div id="msg"></div>
    <input id="input" placeholder="输入消息..." onkeypress="if(event.key==='Enter')send()"><button onclick="send()">发送</button>
    <script>
        let ws;
        const status=document.getElementById('status'),msgDiv=document.getElementById('msg');
        function connect(){
            ws=new WebSocket(`${location.protocol==='https:'?'wss:':'ws:'}//${location.host}/ws`);
            ws.onopen=()=>{status.innerHTML='<span style=color:green>已连接</span>'};
            ws.onmessage=(e)=>{const d=JSON.parse(e.data);msgDiv.innerHTML+=`<div>[${new Date().toLocaleTimeString()}] ${d.message}</div>`};
            ws.onclose=()=>{status.innerHTML='<span style=color:red>已断开</span>';setTimeout(connect,3000)};
        }
        function send(){
            if(!ws||ws.readyState!==WebSocket.OPEN)return;
            ws.send(JSON.stringify({content:document.getElementById('input').value}));
            document.getElementById('input').value='';
        }
        window.onload=connect;
    </script>
</body>
</html>
"""

@app.get("/")
async def get_page():
    return HTMLResponse(html)

@app.websocket("/ws")
async def ws_endpoint(websocket: WebSocket):
    await websocket.accept()
    try:
        while True:
            data = await websocket.receive_text()
            content = json.loads(data).get("content", data)
            await websocket.send_json({
                "message": f"服务器收到：{content}",
                "time": datetime.utcnow().isoformat()
            })
    except WebSocketDisconnect:
        print("连接断开")

连接管理器实现

单连接没用，企业级需要多用户管理、房间隔离、系统通知，用一个连接管理器统一调度：

from fastapi import WebSocket, WebSocketDisconnect
from typing import Dict, Set
import json
from datetime import datetime

class ConnectionManager:
    def __init__(self):
        # 全局活跃连接：user_id -> WebSocket
        self.active_connections: Dict[str, WebSocket] = {}
        # 房间管理：room_id -> Set(user_id)
        self.room_connections: Dict[str, Set[str]] = {}

    async def connect(self, user_id: str, websocket: WebSocket):
        await websocket.accept()
        self.active_connections[user_id] = websocket
        await self.broadcast_system(f"用户 {user_id} 上线")

    def disconnect(self, user_id: str):
        if user_id in self.active_connections:
            del self.active_connections[user_id]
            for room in self.room_connections.values():
                room.discard(user_id)
            # 清理空房间
            self.room_connections = {k:v for k,v in self.room_connections.items() if v}

    async def send_personal(self, user_id: str, message: dict):
        if user_id in self.active_connections:
            try:
                await self.active_connections[user_id].send_json(message)
            except Exception:
                self.disconnect(user_id)

    async def broadcast_room(self, room_id: str, message: dict, exclude_user: str = None):
        if room_id not in self.room_connections:
            return
        for user_id in self.room_connections[room_id]:
            if exclude_user and user_id == exclude_user:
                continue
            await self.send_personal(user_id, message)

    async def broadcast_system(self, content: str):
        message = {
            "type": "system",
            "content": content,
            "time": datetime.utcnow().isoformat()
        }
        for user_id in list(self.active_connections.keys()):
            await self.send_personal(user_id, message)

    async def join_room(self, user_id: str, room_id: str):
        if room_id not in self.room_connections:
            self.room_connections[room_id] = set()
        self.room_connections[room_id].add(user_id)
        await self.broadcast_room(room_id, {
            "type": "join",
            "content": f"用户 {user_id} 加入房间",
            "time": datetime.utcnow().isoformat()
        }, exclude_user=user_id)

# 全局实例
manager = ConnectionManager()

# 带连接管理的端点
@app.websocket("/ws/chat/{user_id}/{room_id}")
async def chat_ws(websocket: WebSocket, user_id: str, room_id: str):
    await manager.connect(user_id, websocket)
    await manager.join_room(user_id, room_id)
    try:
        while True:
            data = await websocket.receive_text()
            content = json.loads(data).get("content", data)
            await manager.broadcast_room(room_id, {
                "type": "chat",
                "user_id": user_id,
                "content": content,
                "time": datetime.utcnow().isoformat()
            })
    except WebSocketDisconnect:
        manager.disconnect(user_id)
        await manager.broadcast_system(f"用户 {user_id} 下线")

AI实时助手集成

结合OpenAI的流式API，实现打字机效果的AI对话，这里用FastAPI的异步特性：

from openai import AsyncOpenAI
import os

# 初始化AI客户端（实际用环境变量存储key）
ai_client = AsyncOpenAI(api_key=os.getenv("OPENAI_API_KEY"))

@app.websocket("/ws/ai/{user_id}")
async def ai_ws(websocket: WebSocket, user_id: str):
    await websocket.accept()
    try:
        while True:
            data = await websocket.receive_text()
            user_msg = json.loads(data).get("content", "")
            if not user_msg:
                continue
            
            # 启动流式调用
            response = await ai_client.chat.completions.create(
                model="gpt-4o-mini",
                messages=[{"role": "user", "content": user_msg}],
                stream=True,
                temperature=0.7
            )
            
            # 实时发送token
            full_resp = ""
            async for chunk in response:
                if chunk.choices[0].delta.content:
                    token = chunk.choices[0].delta.content
                    full_resp += token
                    await websocket.send_json({
                        "type": "ai_token",
                        "token": token,
                        "partial": full_resp
                    })
            # 发送完成消息
            await websocket.send_json({
                "type": "ai_done",
                "full": full_resp
            })
    except WebSocketDisconnect:
        print(f"AI会话 {user_id} 断开")

核心最佳实践

1. 安全认证

不要在URL明文带token，用WebSocket Query参数+JWT验证，并限制速率和并发：

from fastapi import Query
from jose import JWTError, jwt
from auth.config import SECRET_KEY, ALGORITHM  # 替换为你的JWT配置

# 连接管理器加上认证
async def verify_ws_token(token: str = Query(...)):
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        return payload.get("sub")  # 返回user_id
    except JWTError:
        return None

# 改造端点
@app.websocket("/ws/secure/chat/{room_id}")
async def secure_ws(websocket: WebSocket, room_id: str, user_id: str = Depends(verify_ws_token)):
    if not user_id:
        await websocket.close(code=4001, reason="认证失败")
        return
    # 后续逻辑同前

2. 性能优化

  • 消息批处理：用asyncio.Queue收集高频消息，批量发送
  • 连接超时清理：定期删除无活动的连接（用Redis的TTL更好）
  • 禁用压缩：频繁小消息压缩反而慢，大消息再开

3. 集群部署

用Redis发布订阅实现跨节点消息传递，负载均衡器选支持WebSocket的（Nginx、HAProxy、Cloudflare）：

# Nginx反向代理WebSocket配置
location /ws/ {
    proxy_pass http://backend_fastapi;
    proxy_http_version 1.1;
    proxy_set_header Upgrade $http_upgrade;
    proxy_set_header Connection "upgrade";
    proxy_set_header Host $host;
    proxy_set_header X-Real-IP $remote_addr;
    proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
    proxy_read_timeout 86400;  # 延长超时
}

总结

FastAPI原生支持WebSocket，配合异步特性能轻松构建高性能实时应用：

  1. 用连接管理器统一管理多用户、房间
  2. 集成JWT、速率限制保证安全
  3. 用Redis+负载均衡实现集群扩展
  4. 结合AI流式API打造真实的实时助手

💡 实战建议：从小规模开始，先测试单连接→单房间→多房间→集群，逐步优化性能和稳定性。