102 lines
3.4 KiB
Python
102 lines
3.4 KiB
Python
"""WebSocket 路由"""
|
|
|
|
import json
|
|
|
|
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Depends, Query
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.database import async_session
|
|
from app.utils.security import decode_access_token
|
|
from app.websocket.manager import manager
|
|
from app.websocket.events import EventType
|
|
from app.websocket.handlers import (
|
|
handle_chat_send, handle_chat_typing,
|
|
handle_chat_read, handle_presence_update,
|
|
)
|
|
|
|
websocket_router = APIRouter()
|
|
|
|
|
|
@websocket_router.websocket("/ws")
|
|
async def websocket_endpoint(websocket: WebSocket, token: str = Query(None)):
|
|
"""WebSocket 连接端点"""
|
|
|
|
# 验证 Token
|
|
if not token:
|
|
await websocket.close(code=4001, reason="Missing token")
|
|
return
|
|
|
|
payload = decode_access_token(token)
|
|
if not payload:
|
|
await websocket.close(code=4001, reason="Invalid token")
|
|
return
|
|
|
|
user_id = payload.get("sub")
|
|
if not user_id:
|
|
await websocket.close(code=4001, reason="Invalid token payload")
|
|
return
|
|
|
|
# 接受连接
|
|
await manager.connect(websocket, user_id)
|
|
|
|
# 更新在线状态
|
|
async with async_session() as db:
|
|
from app.services.user_service import UserService
|
|
user_service = UserService(db)
|
|
await user_service.update_status(user_id, "online")
|
|
await db.commit()
|
|
|
|
# 广播上线通知
|
|
await manager.broadcast(EventType.PRESENCE_ONLINE, {"user_id": user_id})
|
|
print(f"🌿 用户 {user_id} 已连接 WebSocket")
|
|
|
|
try:
|
|
while True:
|
|
raw = await websocket.receive_text()
|
|
try:
|
|
event = json.loads(raw)
|
|
except json.JSONDecodeError:
|
|
await manager.send_to_user(user_id, EventType.ERROR, {"message": "无效的 JSON"})
|
|
continue
|
|
|
|
event_type = event.get("type")
|
|
data = event.get("data", {})
|
|
|
|
# 创建新的数据库会话处理事件
|
|
async with async_session() as db:
|
|
handler_map = {
|
|
EventType.CHAT_SEND: handle_chat_send,
|
|
EventType.CHAT_TYPING: handle_chat_typing,
|
|
EventType.CHAT_READ: handle_chat_read,
|
|
EventType.PRESENCE_UPDATE: handle_presence_update,
|
|
}
|
|
|
|
handler = handler_map.get(event_type)
|
|
if handler:
|
|
try:
|
|
await handler(websocket, user_id, data, db)
|
|
await db.commit()
|
|
except Exception as e:
|
|
await manager.send_to_user(
|
|
user_id, EventType.ERROR, {"message": str(e)}
|
|
)
|
|
await db.rollback()
|
|
else:
|
|
await manager.send_to_user(
|
|
user_id, EventType.ERROR,
|
|
{"message": f"未知事件类型: {event_type}"}
|
|
)
|
|
|
|
except WebSocketDisconnect:
|
|
manager.disconnect(websocket, user_id)
|
|
|
|
# 更新离线状态
|
|
async with async_session() as db:
|
|
from app.services.user_service import UserService
|
|
user_service = UserService(db)
|
|
await user_service.update_status(user_id, "offline")
|
|
await db.commit()
|
|
|
|
await manager.broadcast(EventType.PRESENCE_OFFLINE, {"user_id": user_id})
|
|
print(f"🌿 用户 {user_id} 已断开 WebSocket")
|