# -*- coding: utf-8 -*- """ 聊天 WebSocket 消费者 """ import json import logging from typing import Dict, Any from fastapi import WebSocket from core.websocket.consumers.base import TokenAuthWebSocketConsumer, manager logger = logging.getLogger(__name__) class ChatConsumer(TokenAuthWebSocketConsumer): """聊天WebSocket消费者""" def __init__(self, websocket: WebSocket): super().__init__(websocket) self._conversation_ids: set = set() async def connect(self): """连接并加入用户聊天组""" await super().connect() if self.is_authenticated and self.user_id: # 加入用户聊天组（用于接收所有聊天消息） await manager.group_add( f"chat_user_{self.user_id}", self.websocket, ) # 广播上线事件给相关会话成员 await self._broadcast_presence("online") async def disconnect(self, close_code: int = 1000): """断开连接""" if self.user_id: # 先广播下线事件（在移除连接之前检查是否还有其他连接） await manager.group_discard( f"chat_user_{self.user_id}", self.websocket, ) # 离开所有会话组 for conv_id in self._conversation_ids: await manager.group_discard( f"chat_conv_{conv_id}", self.websocket, ) await super().disconnect(close_code) # 断开后检查用户是否还有其他连接，没有则广播离线 if self.user_id and not manager.is_online(self.user_id): await self._broadcast_presence("offline") async def _broadcast_presence(self, status: str): """广播用户在线状态变更给所有相关会话成员""" try: from app.database import AsyncSessionLocal from core.chat.service import ConversationService async with AsyncSessionLocal() as db: # 获取该用户参与的所有会话的成员 related_user_ids = await ConversationService.get_related_user_ids(db, self.user_id) presence_data = { "type": "chat.presence", "data": { "user_id": self.user_id, "status": status, }, } for uid in related_user_ids: if uid != self.user_id: await manager.send_to_user(uid, presence_data) except Exception as e: logger.error(f"广播在线状态失败: {e}") async def handle_message(self, data: Dict[str, Any]): """处理聊天相关消息""" message_type = data.get("type", "unknown") payload = data.get("data", {}) if message_type == "chat.send": await self._handle_send(payload) elif message_type == "chat.typing": await self._handle_typing(payload) elif message_type == "chat.read": await self._handle_read(payload) elif message_type == "chat.join": await self._handle_join(payload) else: await self.send_error(f"未知消息类型: {message_type}") async def _handle_send(self, payload: Dict[str, Any]): """处理发送消息""" conversation_id = payload.get("conversation_id") msg_type = payload.get("msg_type", "text") content = payload.get("content") file_id = payload.get("file_id") reply_to_id = payload.get("reply_to_id") extra = payload.get("extra") if not conversation_id: await self.send_error("缺少 conversation_id") return try: from app.database import AsyncSessionLocal from core.chat.service import ConversationService, ChatMessageService async with AsyncSessionLocal() as db: # 验证成员身份 if not await ConversationService.is_member(db, conversation_id, self.user_id): await self.send_error("非会话成员") return # 发送消息 msg = await ChatMessageService.send_message( db, conversation_id=conversation_id, sender_id=self.user_id, msg_type=msg_type, content=content, file_id=file_id, reply_to_id=reply_to_id, extra=extra, ) # 获取发送者信息 from core.user.model import User from sqlalchemy import select sender_result = await db.execute(select(User).where(User.id == self.user_id)) sender = sender_result.scalar_one_or_none() # 获取文件信息 file_info = {} if msg.file_id: from core.file_manager.model import FileManager file_result = await db.execute( select(FileManager).where(FileManager.id == msg.file_id) ) f = file_result.scalar_one_or_none() if f: file_info = { "file_name": f.name, "file_url": f"/api/core/file_manager/stream/{f.id}", "file_size": f.size, "file_ext": f.file_ext, } # 推送给会话所有成员 member_ids = await ConversationService.get_member_user_ids(db, conversation_id) message_data = { "type": "chat.message", "data": { "id": msg.id, "conversation_id": msg.conversation_id, "sender_id": msg.sender_id, "msg_type": msg.msg_type, "content": msg.content, "file_id": msg.file_id, "reply_to_id": msg.reply_to_id, "is_recalled": msg.is_recalled, "extra": msg.extra, "sys_create_datetime": msg.sys_create_datetime.isoformat() if msg.sys_create_datetime else None, "sender_name": (sender.name or sender.username) if sender else None, "sender_avatar": sender.avatar if sender else None, **file_info, }, } for uid in member_ids: await manager.send_to_user(uid, message_data) except Exception as e: logger.error(f"发送消息失败: {e}") await self.send_error(f"发送消息失败: {str(e)}") async def _handle_typing(self, payload: Dict[str, Any]): """处理正在输入状态""" conversation_id = payload.get("conversation_id") if not conversation_id: return try: from app.database import AsyncSessionLocal from core.chat.service import ConversationService from core.user.model import User from sqlalchemy import select async with AsyncSessionLocal() as db: member_ids = await ConversationService.get_member_user_ids(db, conversation_id) sender_result = await db.execute(select(User).where(User.id == self.user_id)) sender = sender_result.scalar_one_or_none() typing_data = { "type": "chat.typing", "data": { "conversation_id": conversation_id, "user_id": self.user_id, "user_name": (sender.name or sender.username) if sender else None, }, } for uid in member_ids: if uid != self.user_id: await manager.send_to_user(uid, typing_data) except Exception as e: logger.error(f"发送typing状态失败: {e}") async def _handle_read(self, payload: Dict[str, Any]): """处理已读标记""" conversation_id = payload.get("conversation_id") message_id = payload.get("message_id") if not conversation_id or not message_id: return try: from app.database import AsyncSessionLocal from core.chat.service import ConversationService, ChatMessageService async with AsyncSessionLocal() as db: await ChatMessageService.mark_read(db, conversation_id, self.user_id, message_id) # 通知其他成员已读回执 member_ids = await ConversationService.get_member_user_ids(db, conversation_id) read_data = { "type": "chat.read_receipt", "data": { "conversation_id": conversation_id, "user_id": self.user_id, "message_id": message_id, }, } for uid in member_ids: if uid != self.user_id: await manager.send_to_user(uid, read_data) except Exception as e: logger.error(f"标记已读失败: {e}") async def _handle_join(self, payload: Dict[str, Any]): """加入会话组（用于接收该会话的实时消息）""" conversation_id = payload.get("conversation_id") if not conversation_id: return self._conversation_ids.add(conversation_id) await manager.group_add(f"chat_conv_{conversation_id}", self.websocket) await self.send_message("chat.joined", f"已加入会话 {conversation_id}")