fastapi websocket全双工通信

下面是一个基于 FastAPI WebSocket 的完整实现，核心思路是用 asyncio.Task 管理每次回复任务，收到新问题时通过 task.cancel() 终止上一个回复，再启动新的回复。

架构设计

客户端 ──WebSocket──▶ 服务端
                        │
                        ├── 主循环：持续接收客户端消息
                        │       │
                        │       ├── 收到新问题 → cancel() 旧任务
                        │       │               → 创建新 Task
                        │       │
                        │       └── 收到 cancel 指令 → cancel() 当前任务
                        │
                        └── 回复协程：流式发送回复内容
                                （可被 cancel 中断）

完整代码

import asyncio
import json
import uuid
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse

app = FastAPI()


# ========== 模拟 AI 流式回复 ==========
async def fake_stream_response(query: str, websocket: WebSocket, stop_event: asyncio.Event):
    """
    模拟一个流式 AI 回复。
    每 0.1 秒发送一个 token，如果 stop_event 被设置，则提前终止。
    """
    fake_tokens = [
        f"针对您的问题「{query}」，",
        "我来详细解答一下。",
        "首先，",
        "这是一个非常好的问题。",
        "从技术角度来看，",
        "我们需要考虑以下几个方面：",
        "第一，性能优化；",
        "第二，代码可读性；",
        "第三，系统可扩展性。",
        "综上所述，",
        "建议您采用异步架构。",
        "希望这个回答对您有帮助！"
    ]

    for i, token in enumerate(fake_tokens):
        if stop_event.is_set():
            # 被取消，发送中断通知
            await websocket.send_json({
                "type": "response.interrupted",
                "message": "回复已被新请求中断"
            })
            return

        await websocket.send_json({
            "type": "response.token",
            "token": token,
            "index": i
        })
        await asyncio.sleep(0.3)  # 模拟生成延迟

    # 回复完成
    await websocket.send_json({
        "type": "response.done",
        "message": "回复完成"
    })


# ========== WebSocket 端点 ==========
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()

    # 当前正在执行的回复任务
    current_task: asyncio.Task | None = None
    # 用于通知回复协程停止的信号
    stop_event = asyncio.Event()

    try:
        while True:
            # 持续监听客户端消息
            raw_data = await websocket.receive_text()
            data = json.loads(raw_data)
            msg_type = data.get("type")

            # ---------- 收到新问题 ----------
            if msg_type == "query":
                query_text = data.get("content", "")
                query_id = str(uuid.uuid4())[:8]

                # 1. 如果有正在进行的回复，先取消它
                if current_task and not current_task.done():
                    stop_event.set()          # 通知旧协程停止
                    current_task.cancel()     # 取消旧任务
                    try:
                        await current_task   # 等待旧任务清理
                    except asyncio.CancelledError:
                        pass

                # 2. 重置 stop_event，创建新的回复任务
                stop_event.clear()
                current_task = asyncio.create_task(
                    fake_stream_response(query_text, websocket, stop_event)
                )

                # 通知客户端：新回复已开始
                await websocket.send_json({
                    "type": "response.started",
                    "query_id": query_id,
                    "query": query_text
                })

            # ---------- 客户端主动取消 ----------
            elif msg_type == "cancel":
                if current_task and not current_task.done():
                    stop_event.set()
                    current_task.cancel()
                    try:
                        await current_task
                    except asyncio.CancelledError:
                        pass
                    await websocket.send_json({
                        "type": "response.interrupted",
                        "message": "用户主动取消"
                    })

    except WebSocketDisconnect:
        # 客户端断开，清理资源
        if current_task and not current_task.done():
            current_task.cancel()
        print("客户端已断开连接")


# ========== 测试页面 ==========
@app.get("/")
async def index():
    return HTMLResponse(content=TEST_HTML)


TEST_HTML = """
<!DOCTYPE html>
<html lang="zh-CN">
<head>
    <meta charset="UTF-8">
    <title>全双工 WebSocket 测试</title>
    <style>
        * { margin: 0; padding: 0; box-sizing: border-box; }
        body { font-family: -apple-system, sans-serif; background: #f5f5f5; padding: 20px; }
        .container { max-width: 700px; margin: 0 auto; }
        h1 { text-align: center; margin-bottom: 20px; color: #333; }
        .chat-box {
            background: #fff; border-radius: 12px; padding: 20px;
            height: 450px; overflow-y: auto; margin-bottom: 16px;
            box-shadow: 0 2px 12px rgba(0,0,0,0.08);
        }
        .msg { margin: 8px 0; padding: 10px 14px; border-radius: 8px; max-width: 85%; }
        .msg.user { background: #007aff; color: #fff; margin-left: auto; text-align: right; }
        .msg.server { background: #e9e9eb; color: #333; }
        .msg.system { background: #fff3cd; color: #856404; font-size: 13px; text-align: center; max-width: 100%; }
        .input-area { display: flex; gap: 10px; }
        input { flex: 1; padding: 12px; border: 1px solid #ddd; border-radius: 8px; font-size: 15px; }
        button {
            padding: 12px 20px; border: none; border-radius: 8px;
            cursor: pointer; font-size: 15px; color: #fff;
        }
        .btn-send { background: #007aff; }
        .btn-cancel { background: #ff3b30; }
        button:disabled { opacity: 0.5; cursor: not-allowed; }
        .status { text-align: center; margin: 10px 0; font-size: 13px; color: #999; }
    </style>
</head>
<body>
    <div class="container">
        <h1>🔄 全双工 WebSocket 演示</h1>
        <div class="status" id="status">未连接</div>
        <div class="chat-box" id="chatBox"></div>
        <div class="input-area">
            <input id="msgInput" placeholder="输入问题，发送可中断上一个回复..." 
                   onkeydown="if(event.key==='Enter') sendQuery()">
            <button class="btn-send" onclick="sendQuery()">发送</button>
            <button class="btn-cancel" id="cancelBtn" onclick="sendCancel()" disabled>中断</button>
        </div>
    </div>

    <script>
        const chatBox = document.getElementById('chatBox');
        const msgInput = document.getElementById('msgInput');
        const cancelBtn = document.getElementById('cancelBtn');
        const statusEl = document.getElementById('status');
        let ws;
        let currentResponseEl = null; // 当前正在接收的回复元素

        function connect() {
            ws = new WebSocket('ws://' + location.host + '/ws');

            ws.onopen = () => {
                statusEl.textContent = '✅ 已连接';
                statusEl.style.color = '#34c759';
                addMessage('system', '连接成功！试试快速发送多个问题。');
            };

            ws.onmessage = (event) => {
                const data = JSON.parse(event.data);

                if (data.type === 'response.started') {
                    // 新回复开始，创建新的消息气泡
                    currentResponseEl = document.createElement('div');
                    currentResponseEl.className = 'msg server';
                    currentResponseEl.textContent = '';
                    chatBox.appendChild(currentResponseEl);
                    cancelBtn.disabled = false;
                }
                else if (data.type === 'response.token') {
                    // 追加 token 到当前回复
                    if (currentResponseEl) {
                        currentResponseEl.textContent += data.token;
                    }
                }
                else if (data.type === 'response.done') {
                    currentResponseEl = null;
                    cancelBtn.disabled = true;
                }
                else if (data.type === 'response.interrupted') {
                    if (currentResponseEl) {
                        currentResponseEl.textContent += '\\n⚠️ ' + data.message;
                        currentResponseEl.style.borderLeft = '3px solid #ff9500';
                    }
                    currentResponseEl = null;
                    cancelBtn.disabled = true;
                }

                chatBox.scrollTop = chatBox.scrollHeight;
            };

            ws.onclose = () => {
                statusEl.textContent = '❌ 连接断开，3秒后重连...';
                statusEl.style.color = '#ff3b30';
                setTimeout(connect, 3000);
            };
        }

        function sendQuery() {
            const text = msgInput.value.trim();
            if (!text || !ws || ws.readyState !== WebSocket.OPEN) return;

            addMessage('user', text);
            ws.send(JSON.stringify({ type: 'query', content: text }));
            msgInput.value = '';
        }

        function sendCancel() {
            if (ws && ws.readyState === WebSocket.OPEN) {
                ws.send(JSON.stringify({ type: 'cancel' }));
            }
        }

        function addMessage(role, text) {
            const el = document.createElement('div');
            el.className = 'msg ' + role;
            el.textContent = text;
            chatBox.appendChild(el);
            chatBox.scrollTop = chatBox.scrollHeight;
        }

        connect();
    </script>
</body>
</html>
"""

启动方式

pip install fastapi uvicorn
uvicorn main:app --reload --port 8000

然后浏览器打开 http://localhost:8000 即可测试。

核心机制解析

整个"中断旧回复、响应新问题"的流程，关键在于三个要素的配合：

要素	作用	对应代码
asyncio.Task	将每次回复包装为独立任务，可被外部取消	asyncio.create_task(...)
task.cancel()	向任务抛出 CancelledError，终止其执行	收到新 query 时调用
asyncio.Event	协作式停止信号，让回复协程在循环中主动检查并退出	stop_event.is_set()

执行流程：

1. 用户发送问题 A
   → 创建 Task-A，开始流式回复 A...

2. 用户发送问题 B（此时 A 还在回复中）
   → stop_event.set()     # 通知 Task-A 停止
   → Task-A.cancel()      # 取消 Task-A
   → await Task-A         # 等待清理完成
   → stop_event.clear()   # 重置信号
   → 创建 Task-B，开始流式回复 B...

3. Task-A 在下一个 token 发送前检测到 stop_event
   → 发送 "response.interrupted" 通知客户端
   → 退出协程

协议约定

客户端和服务端通过 JSON 消息中的 type 字段区分事件类型：

客户端 → 服务端：

{"type": "query", "content": "用户的问题"}
{"type": "cancel"}

服务端 → 客户端：

{"type": "response.started", "query_id": "a1b2c3d4", "query": "..."}
{"type": "response.token", "token": "片段内容", "index": 0}
{"type": "response.done"}
{"type": "response.interrupted", "message": "回复已被新请求中断"}

这个架构可以直接对接真实的 LLM API（如 OpenAI、Claude），只需将 fake_stream_response 替换为真实的流式 API 调用即可。