Files
wecom_it_smart_desk/backend/app/api/ws.py
T

590 lines
26 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
# =============================================================================
# 企微IT智能服务台 — WebSocket 端点
# =============================================================================
# 说明:提供 WebSocket 端点,供坐席前端和H5用户端建立长连接,实现实时推送。
# 核心功能:
# 1. 接受坐席的 WebSocket 连接请求(含 token 认证)— /ws/{agent_id}
# 2. 接受H5员工的 WebSocket 连接请求(含 token 认证)— /ws/h5/{employee_id}
# 3. 维持连接,监听客户端消息(主要是心跳 ping)
# 4. 连接断开时自动清理注册信息
# 安全(WS-01):
# 握手时从 query param 取 token → 查 Redis 验证 → 不通过则 close(code=4001)
# 防止未授权用户冒充坐席/员工建立 WS 连接
#
# 端点路径:
# - 坐席端:/ws/{agent_id}?token=xxx
# - H5员工端:/ws/h5/{employee_id}?token=xxx
# 为什么不挂 /api 前缀:WebSocket 不是 REST API,不走 Vite 的 /api 代理配置
# =============================================================================
import logging
import asyncio
from fastapi import APIRouter, WebSocket, WebSocketDisconnect
from app.services.ws_manager import manager as ws_manager
from app.services.cache_service import cache_service
from app.database import _get_session_factory
from app.models.conversation import Conversation
from app.models.message import Message
from app.tasks.h5_ai_task import process_h5_ai_reply
logger = logging.getLogger(__name__)
# WebSocket 路由器(不挂 /api 前缀,直接注册在应用根路径)
router = APIRouter()
# 认证失败时的 WebSocket 关闭码
# 4001 = 自定义码,表示"未授权"(4000+ 为应用自定义范围)
WS_CLOSE_UNAUTHORIZED = 4001
@router.websocket("/ws/{agent_id}")
async def websocket_endpoint(
websocket: WebSocket,
agent_id: str,
) -> None:
"""坐席 WebSocket 端点主循环(含 WS-01 token 认证)。
做什么:
1. 从 Authorization header 获取 token(优先)或 query param(兼容)
2. 验证 token 有效性(查 Redis
3. 验证 token 与 agent_id 一致性(防冒充)
4. 认证通过后接受连接,注册到 ConnectionManager
5. 进入消息接收循环,处理客户端发送的消息
6. 连接断开时清理注册信息
为什么需要 token 认证(WS-01):
- 之前 /ws/{agent_id} 无任何认证,任何人知道 URL 即可冒充任意坐席
- 攻击者可监听所有消息、发送伪造消息,是 P0 级安全漏洞
- 修复后,必须提供与 agent_id 匹配的有效 token 才能建立连接
安全改进(P0-#4):
- 优先从 Authorization: Bearer {token} header 获取 token
- 兼容从 ?token= URL 参数获取(向后兼容)
- 不再将 token 暴露在 URL 中,避免 access_log 泄露
v0.5.1 修复:移除 `request: Request` 参数(部分 Starlette 版本注入 Request 失败,
改用 `websocket.headers` 和 `websocket.query_params` 读取 header/query)
Args:
websocket: FastAPI WebSocket 对象(框架自动注入)
agent_id: 坐席ID(从 URL 路径参数获取)
"""
# ======================================================================
# WS-01: Token 认证(从 subprotocol / header / query 获取)
# ======================================================================
# 步骤1: 优先从 Sec-WebSocket-Protocol (subprotocol) 获取 token,其次从 Authorization header,最后从 query(向后兼容)
# 格式: Sec-WebSocket-Protocol: bearer.{token}
# 说明: 浏览器原生 WebSocket API 不支持 headers 参数,但支持 subprotocols (第2参数数组)
# 前端用 new WebSocket(url, ["bearer.{token}"]) 传递,服务端从 sec-websocket-protocol 头读取
subprotocol = websocket.headers.get("sec-websocket-protocol", "")
if subprotocol.startswith("bearer."):
token = subprotocol[7:] # 去掉 "bearer." 前缀
else:
# 其次从 Authorization header 获取
auth_header = websocket.headers.get("Authorization", "")
if auth_header.startswith("Bearer "):
token = auth_header[7:] # 去掉 "Bearer " 前缀
else:
# 向后兼容:从 query param 获取(即将废弃)
token = websocket.query_params.get("token", "")
# 步骤2: 检查 token 是否为空
if not token:
# 先 accept 再 close,否则客户端收不到关闭帧
await websocket.accept()
await websocket.close(code=WS_CLOSE_UNAUTHORIZED, reason="Missing token")
logger.warning(f"WebSocket 拒绝连接: agent_id={agent_id}, 原因=缺少token")
return
# 步骤3: 从 Redis 查询 token 对应的坐席信息
# 支持两种格式:
# 1. 新格式: user:token:{token} -> JSON {employee_id, roles, ...}
# 2. 旧格式: agent:token:{token} -> employee_id
try:
import json
# 先尝试新格式
user_info = await cache_service.get(f"user:token:{token}")
if user_info:
try:
user_data = json.loads(user_info)
stored_agent_id = user_data.get("employee_id")
roles = user_data.get("roles", [])
if "agent" not in roles:
# token 没有 agent 角色
await websocket.accept()
await websocket.close(code=WS_CLOSE_UNAUTHORIZED, reason="No agent role")
return
except (json.JSONDecodeError, TypeError):
stored_agent_id = None
else:
# 兼容旧格式
stored_agent_id = await cache_service.get(f"agent:token:{token}")
except Exception as e:
# Redis 不可用时必须拒绝连接:token 验证依赖 Redis,无法验证身份
# 如果降级放行,攻击者可在 Redis 故障时用任意 agent_id 冒充坐席
logger.error(f"Redis 查询失败,拒绝 WS 连接: agent_id={agent_id}, error={e}")
await websocket.accept()
await websocket.close(
code=WS_CLOSE_UNAUTHORIZED,
reason="Authentication service unavailable"
)
return
# 步骤4: 验证 token 与 agent_id 一致性
if not stored_agent_id:
# token 不存在(已过期或伪造)
await websocket.accept()
await websocket.close(code=WS_CLOSE_UNAUTHORIZED, reason="Invalid or expired token")
logger.warning(f"WebSocket 拒绝连接: agent_id={agent_id}, 原因=token无效或已过期")
return
if stored_agent_id != agent_id:
# token 对应的坐席与请求的 agent_id 不匹配(冒充)
await websocket.accept()
await websocket.close(code=WS_CLOSE_UNAUTHORIZED, reason="Token-agent mismatch")
logger.warning(
f"WebSocket 拒绝连接: agent_id={agent_id}, "
f"原因=token对应坐席{stored_agent_id}与请求不匹配"
)
return
# ======================================================================
# 认证通过,建立连接
# ======================================================================
# 注册连接(内部会调用 websocket.accept(),并回显协商的 subprotocol
await ws_manager.connect(agent_id, websocket, subprotocol=subprotocol)
logger.info(f"坐席 WebSocket 连接已认证: agent_id={agent_id}")
try:
# 消息接收循环
# 保持连接打开,监听客户端发来的消息
# 即使客户端不发消息,这个循环也必须保持,否则连接会关闭
while True:
# 等待接收客户端消息(阻塞等待)
data = await websocket.receive_json()
# 处理心跳 ping
# 前端每 30 秒发送一次 ping,后端回复 pong
# 作用:检测连接是否存活,防止中间代理(如 Nginx)因超时断开连接
if data.get("type") == "ping":
await websocket.send_json({"type": "pong"})
logger.debug(f"WebSocket 心跳: agent_id={agent_id}")
# 处理输入指示器 typing 事件
# 前端在用户输入时发送 typing 事件,后端广播给同一会话的其他参与者
elif data.get("type") == "typing":
conversation_id = data.get("conversation_id")
sender_name = data.get("sender_name", agent_id)
if conversation_id:
# 广播给所有坐席(包含 sender_type 和 sender_id
# 前端可据此过滤掉自己的 typing 事件)
await ws_manager.broadcast({
"type": "typing",
"data": {
"conversation_id": conversation_id,
"sender_id": agent_id,
"sender_name": sender_name,
"sender_type": "agent",
}
})
else:
# 未来可扩展处理其他类型的客户端消息
logger.debug(
f"WebSocket 收到未知消息: agent_id={agent_id}, "
f"type={data.get('type', 'unknown')}"
)
except WebSocketDisconnect:
# 客户端主动断开连接(正常行为)
# 清理 ConnectionManager 中的注册信息
ws_manager.disconnect(agent_id)
logger.info(f"坐席断开 WebSocket 连接: agent_id={agent_id}")
except Exception as e:
# 其他异常(如网络错误、JSON 解析错误等)
# 确保注册信息被清理
ws_manager.disconnect(agent_id)
logger.warning(f"WebSocket 异常断开: agent_id={agent_id}, error={e}")
# ==========================================================================
# H5员工 WebSocket 端点
# ==========================================================================
async def _handle_option_select(
conversation_id: str,
employee_id: str,
option_label: str,
):
"""处理员工点击 AI 选项按钮的后端逻辑(v2.0 新增)。
做什么:
1. 在 DB 中存储员工的选项选择为一条 employee 消息
2. 广播该消息给坐席端(让坐席看到员工选了什么)
3. 触发 process_h5_ai_reply() → Dify 接收选项文本作为用户消息 → 返回下一轮 AI 回复
为什么:前端 sendOptionSelect() 通过 WS 发送 option_select 消息,
后端必须接收并触发 AI 回复,否则用户点击选项后无响应。
Args:
conversation_id: 会话 ID
employee_id: 员工企微 UserID
option_label: 选项的显示文本(如"企微密码"),作为用户消息发给 Dify
"""
from datetime import datetime
try:
async with _get_session_factory()() as db:
# 1. 查找会话(获取 dify_conversation_id 用于多轮上下文)
conversation = await db.get(Conversation, conversation_id)
if not conversation:
logger.warning(f"option_select: 会话不存在 {conversation_id}")
return
# 2. 存储员工消息(选项选择作为文本消息)
emp_msg = Message(
conversation_id=conversation_id,
sender_type="employee",
sender_id=employee_id,
sender_name="", # 前端会从 employeeStore 补全
content=option_label,
msg_type="text",
is_read=False,
)
db.add(emp_msg)
await db.flush()
# 更新会话时间
conversation.updated_at = datetime.now()
db.add(conversation)
await db.flush()
await db.commit()
# 3. 广播给坐席端(让坐席看到员工的选择)
try:
await ws_manager.broadcast({
"type": "new_message",
"data": {
"conversation_id": str(conversation_id),
"message_id": str(emp_msg.id),
"sender_type": "employee",
"sender_id": employee_id,
"content": option_label,
"msg_type": "text",
},
})
except Exception as ws_err:
logger.warning(f"option_select: WS 广播坐席失败: {ws_err}")
# 4. 触发 AI 回复(异步后台任务,不阻塞)
# dify_conversation_id 从 conversation 对象获取(保持多轮上下文)
asyncio.create_task(
process_h5_ai_reply(
conversation_id=conversation_id,
employee_id=employee_id,
content=option_label,
dify_conversation_id=conversation.dify_conversation_id,
msg_type="text",
)
)
logger.info(
f"option_select 已触发 AI 回复: conv={conversation_id}, "
f"option={option_label}"
)
except Exception as e:
logger.error(f"option_select 处理异常: {e}", exc_info=True)
@router.websocket("/ws/h5/{employee_id}")
async def h5_websocket_endpoint(
websocket: WebSocket,
employee_id: str,
) -> None:
"""H5员工 WebSocket 端点主循环(含 token 认证)。
做什么:
1. 从 Authorization header 获取 token(优先从)或 query param(兼容)
2. 验证 employee token 有效性(查 Redis
3. 验证 token 与 employee_id 一致性(防冒充)
4. 认证通过后接受连接,注册到 ConnectionManager 的员工连接表
5. 进入消息接收循环,处理心跳 ping
6. 连接断开时清理注册信息
为什么需要 H5 WS 连接:
- H5员工需要实时接收参与者变更事件(新参与者加入、有人退出等)
- 当前仅通过 3 秒轮询获取更新,实时性不足
- WS 推送 + 轮询降级,双通道保证消息可达
安全改进(P0-#4):
- 优先从 Authorization: Bearer {token} header 获取 token
- 兼容从 ?token= URL 参数获取(向后兼容)
认证机制(与坐席端一致):
- Redis 中存储格式: employee:token:{token} -> employee_id
- (与H5登录 API /api/h5/mock-login 存储格式一致)
- token 缺失、无效、过期、与 employee_id 不匹配均拒绝连接
v0.5.1 修复:移除 `request: Request` 参数(部分 Starlette 版本注入 Request 失败,
改用 `websocket.headers` 和 `websocket.query_params` 读取 header/query)
Args:
websocket: FastAPI WebSocket 对象(框架自动注入)
employee_id: 员工企微 UserID(从 URL 路径参数获取)
"""
# ======================================================================
# Token 认证(从 subprotocol / header / query 获取)
# ======================================================================
# 步骤1: 优先从 Sec-WebSocket-Protocol (subprotocol) 获取 token,其次从 Authorization header,最后从 query(向后兼容)
# 格式: Sec-WebSocket-Protocol: bearer.{token}
subprotocol = websocket.headers.get("sec-websocket-protocol", "")
if subprotocol.startswith("bearer."):
token = subprotocol[7:] # 去掉 "bearer." 前缀
else:
# 其次从 Authorization header 获取
auth_header = websocket.headers.get("Authorization", "")
if auth_header.startswith("Bearer "):
token = auth_header[7:] # 去掉 "Bearer " 前缀
else:
# 向后兼容:从 query param 获取(即将废弃)
token = websocket.query_params.get("token", "")
# 步骤2: 检查 token 是否为空
if not token:
await websocket.accept()
await websocket.close(code=WS_CLOSE_UNAUTHORIZED, reason="Missing token")
logger.warning(f"H5 WebSocket 拒绝连接: employee_id={employee_id}, 原因=缺少token")
return
# 步骤3: 从 Redis 查询 token 对应的员工信息
# Redis 中存储格式: employee:token:{token} -> employee_id
# (与H5登录 API /api/h5/mock-login 存储格式一致)
try:
stored_employee_id = await cache_service.get(f"employee:token:{token}")
except Exception as e:
# Redis 不可用时必须拒绝连接(与坐席端一致的安全策略)
logger.error(f"Redis 查询失败,拒绝 H5 WS 连接: employee_id={employee_id}, error={e}")
await websocket.accept()
await websocket.close(
code=WS_CLOSE_UNAUTHORIZED,
reason="Authentication service unavailable"
)
return
# 步骤4: 验证 token 与 employee_id 一致性
if not stored_employee_id:
await websocket.accept()
await websocket.close(code=WS_CLOSE_UNAUTHORIZED, reason="Invalid or expired token")
logger.warning(f"H5 WebSocket 拒绝连接: employee_id={employee_id}, 原因=token无效或已过期")
return
if stored_employee_id != employee_id:
await websocket.accept()
await websocket.close(code=WS_CLOSE_UNAUTHORIZED, reason="Token-employee mismatch")
logger.warning(
f"H5 WebSocket 拒绝连接: employee_id={employee_id}, "
f"原因=token对应员工{stored_employee_id}与请求不匹配"
)
return
# ======================================================================
# 认证通过,建立连接
# ======================================================================
# 注册员工连接(内部会调用 websocket.accept(),并回显协商的 subprotocol
await ws_manager.connect_employee(employee_id, websocket, subprotocol=subprotocol)
logger.info(f"H5员工 WebSocket 连接已认证: employee_id={employee_id}")
try:
# 消息接收循环
# H5员工端发送心跳 ping 和 option_select(选项按钮点击)
while True:
data = await websocket.receive_json()
# 处理心跳 ping
if data.get("type") == "ping":
await websocket.send_json({"type": "pong"})
logger.debug(f"H5 WebSocket 心跳: employee_id={employee_id}")
# v2.0: 处理选项按钮点击(option_select
# 做什么:员工点击 AI 结构化消息中的选项按钮后,前端通过 WS 发送 option_select
# 后端接收后触发 AI 回复流程(与普通发消息等效),实现交互式排查闭环
elif data.get("type") == "option_select":
option_data = data.get("data", {})
conv_id = option_data.get("conversation_id")
option_label = option_data.get("option_label", "")
option_value = option_data.get("option_value", "")
if conv_id and option_label:
logger.info(
f"H5 WS option_select: employee={employee_id}, "
f"conv={conv_id}, option={option_value}"
)
# 异步触发 AI 回复(不阻塞 WS 循环)
# process_h5_ai_reply 内部创建独立 DB session
# dify_conversation_id 传 None 时会从 conversation 对象回退读取
asyncio.create_task(
_handle_option_select(
conversation_id=conv_id,
employee_id=employee_id,
option_label=option_label,
)
)
else:
logger.warning(
f"H5 WS option_select 数据不完整: employee={employee_id}, "
f"conv_id={conv_id}, label={option_label}"
)
else:
logger.debug(
f"H5 WebSocket 收到未知消息: employee_id={employee_id}, "
f"type={data.get('type', 'unknown')}"
)
except WebSocketDisconnect:
# 客户端主动断开连接
ws_manager.disconnect_employee(employee_id)
logger.info(f"H5员工断开 WebSocket 连接: employee_id={employee_id}")
except Exception as e:
# 其他异常
ws_manager.disconnect_employee(employee_id)
logger.warning(f"H5 WebSocket 异常断开: employee_id={employee_id}, error={e}")
# ==========================================================================
# 终端 WebSocket 端点(小鱼易联终端大屏)
# ==========================================================================
@router.websocket("/ws/terminal/{terminal_sn}")
async def terminal_websocket_endpoint(
websocket: WebSocket,
terminal_sn: str,
) -> None:
"""终端 WebSocket 端点主循环(token可选认证)。
做什么:
1. 从 subprotocol / header / query 获取 token(可选)
2. token 存在时验证有效性,不存在时允许连接(仅接收推送)
3. 注册到 ConnectionManager 的终端连接表
4. 进入消息接收循环,处理心跳 ping 和 request_status
5. 连接断开时清理注册信息
认证策略(与坐席/H5端不同):
- 终端查看状态不需要登录(访客可查看)
- token 可选:无token时仅接收状态推送,不能发预定指令
- 有token时记录登录用户身份(可用于预定操作)
Args:
websocket: FastAPI WebSocket 对象
terminal_sn: 终端序列号(从 URL 路径参数获取)
"""
# ======================================================================
# Token 认证(可选 — 无token也允许连接)
# ======================================================================
# 从 subprotocol / header / query 获取 token
subprotocol = websocket.headers.get("sec-websocket-protocol", "")
if subprotocol.startswith("bearer."):
token = subprotocol[7:]
else:
auth_header = websocket.headers.get("Authorization", "")
if auth_header.startswith("Bearer "):
token = auth_header[7:]
else:
token = websocket.query_params.get("token", "")
# token 可选:无token也允许连接(仅接收推送,不能发预定指令)
# 有token时验证(用于后续预定操作的身份识别)
is_authenticated = False
if token:
try:
import json
# 尝试新格式
user_info = await cache_service.get(f"user:token:{token}")
if user_info:
try:
user_data = json.loads(user_info)
is_authenticated = True
logger.debug(
f"终端 WS 已认证: sn={terminal_sn}, "
f"user={user_data.get('employee_id')}"
)
except (json.JSONDecodeError, TypeError):
pass
else:
# 兼容旧格式
stored_id = await cache_service.get(f"agent:token:{token}")
if stored_id:
is_authenticated = True
except Exception as e:
logger.warning(f"终端 WS token 验证失败(降级为未认证): sn={terminal_sn}, error={e}")
# ======================================================================
# 建立连接(无论是否认证都接受)
# ======================================================================
await ws_manager.connect_terminal(terminal_sn, websocket, subprotocol=subprotocol)
auth_label = "已认证" if is_authenticated else "未认证(仅查看)"
logger.info(f"终端 WebSocket 连接已建立: sn={terminal_sn}, 状态={auth_label}")
try:
# 消息接收循环
while True:
data = await websocket.receive_json()
# 处理心跳 ping
if data.get("type") == "ping":
await websocket.send_json({"type": "pong"})
logger.debug(f"终端 WebSocket 心跳: sn={terminal_sn}")
# 处理状态刷新请求
elif data.get("type") == "request_status":
meetingroom_id = data.get("data", {}).get("meetingroom_id")
if meetingroom_id:
# 获取最新状态并推送
try:
from app.services.meetingroom_service import MeetingroomService
from app.services.wecom_service import WecomService
from app.config import settings
redis_client = settings.create_redis_client()
wecom_service = WecomService(redis_client)
mr_service = MeetingroomService(wecom_service, redis_client)
status_data = await mr_service.get_current_status(meetingroom_id)
await websocket.send_json({
"type": "room_status_update",
"data": {
"meetingroom_id": meetingroom_id,
"status": status_data.get("status"),
"current_meeting": status_data.get("current_meeting"),
"next_meeting": status_data.get("next_meeting"),
"minutes_to_next": status_data.get("minutes_to_next"),
},
})
except Exception as e:
logger.warning(f"终端请求状态刷新失败: sn={terminal_sn}, error={e}")
else:
logger.debug(
f"终端 WebSocket 收到未知消息: sn={terminal_sn}, "
f"type={data.get('type', 'unknown')}"
)
except WebSocketDisconnect:
ws_manager.disconnect_terminal(terminal_sn)
logger.info(f"终端断开 WebSocket 连接: sn={terminal_sn}")
except Exception as e:
ws_manager.disconnect_terminal(terminal_sn)
logger.warning(f"终端 WebSocket 异常断开: sn={terminal_sn}, error={e}")