feat(chat): 完善聊天

core/routers/chat/handlers.py

@@ -10,6 +10,7 @@ from typing import List, Dict
 
 import redis.asyncio as aioredis
 from fastapi import WebSocket
+from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from core.db.crud import CRUD
@@ -26,19 +27,20 @@ socket_manager = SocketManager()
 async def send_new_message(
         websocket: WebSocket,
         db_session: AsyncSession,
-        cache: aioredis.Redis,
         chats: dict,
         current_user: User,
-        incoming_message: schema.Message,
+        incoming_message: schema.WSMessage,
 ):
     chat_guid: str = str(incoming_message.chat_guid)
-    chat: Chat | None = await CRUD(db_session, Chat).get(filters={'chat_guid': chat_guid})
+    chat: Chat | None = await CRUD(db_session, Chat).get(filters={'guid': chat_guid})
+
     # 判断是否为新聊天
     if not chats or chat_guid not in chats:
         if chat:
             chats[chat_guid] = chat.id
             await socket_manager.add_user_to_chat(chat_guid, websocket)
         else:
+            logger.warning(f'chat {chat_guid} not found')
             await socket_manager.send_error('聊天尚未添加', websocket)
             return
 
@@ -49,64 +51,59 @@ async def send_new_message(
             user_id=current_user.user_id,
             chat_id=chat_id,
             message_type=incoming_message.message_type,
-            content=incoming_message.body,
+            content=incoming_message.content,
         )
         db_session.add(message)
-        await db_session.flush()
 
         chat.updated_at = datetime.now()
         db_session.add(chat)
 
+        await db_session.flush()
         await db_session.commit()
-        await db_session.refresh(message, attribute_names=["user", "chat"])
-        await db_session.refresh(chat, attribute_names=["users"])
+        # await db_session.refresh(message, attribute_names=["user", "chat"])
+        # await db_session.refresh(chat, attribute_names=["users"])
 
     except Exception as exc_info:
         await db_session.rollback()
         logger.exception(f"[new_message] Exception, rolling back session, detail: {exc_info}")
         raise exc_info
 
-    send_message_schema = schema.MessageOut(
+    send_message_schema = schema.WSMessageOut(
         type='chat.new_message',
         message_id=message.id,
         message_guid=message.guid,
         message_type=message.message_type,
         chat_guid=chat.guid,
-        user_id=current_user.user_id,
         content=message.content,
-        created_at=message.created_at,
-        friend_user=current_user,
+        friend_user=current_user.dict(),
+    )
+    send_status_replay = schema.WSMessageOut(
+        type='chat.send_status_reply',
+        message_id=message.id,
+        message_guid=message.guid,
+        message_type=message.message_type,
+        chat_guid=chat.guid,
     )
-    await socket_manager.broadcast_to_chat(chat_guid, send_message_schema.dict())
 
-
-@socket_manager.handler('chat.refresh_unread_messages')
-async def refresh_unread_messages(
-        websocket: WebSocket,
-        db_session: AsyncSession,
-        cache: aioredis.Redis,
-        current_user: User,
-        message: schema.Message,
-):
-    pass
+    await socket_manager.broadcast_to_chat(chat_guid, send_message_schema.model_dump_json())
+    await websocket.send_json(send_status_replay.model_dump_json())
 
 
 @socket_manager.handler('chat.message_read')
 async def read_message(
         websocket: WebSocket,
         db_session: AsyncSession,
-        cache: aioredis.Redis,
         current_user: User,
-        message: schema.Message,
+        incoming_message: schema.WSMessage,
 ):
-    last_read_message_id: str | None = message.body
+    last_read_message_id: str | None = incoming_message.content
     if not last_read_message_id:
         return
 
     db_crud = CRUD(db_session, ReadStatus)
     condition = {
         'user_id': current_user.user_id,
-        'chat_guid': message.chat_guid,
+        'chat_guid': incoming_message.chat_guid,
     }
     result = await db_crud.update(filters=condition, data={'last_read_message_id': last_read_message_id})
     if result is None:
@@ -117,3 +114,13 @@ async def read_message(
             },
             filters=condition
         )
+
+
+@socket_manager.handler('chat.ping')
+async def ping(
+        websocket: WebSocket,
+        db_session: AsyncSession,
+        current_user: User,
+        incoming_message: schema.WSMessage,
+):
+    await websocket.send_json({'type': 'chat.send_status_reply', 'status': 'ok'})

core/routers/chat/managers.py

@@ -10,6 +10,7 @@ import logging
 import redis.asyncio as aioredis
 import ujson
 from fastapi import WebSocket
+from fastapi.websockets import WebSocketState
 
 from core.db.engine import get_cache
 
@@ -72,7 +73,7 @@ class SocketManager:
             self.chats[chat_guid] = {websocket}
             await self.pubsub_client.connect()
             pubsub_subscriber = await self.pubsub_client.subscribe(chat_guid)
-            await asyncio.create_task(self._pubsub_data_reader(pubsub_subscriber))
+            asyncio.create_task(self._pubsub_data_reader(pubsub_subscriber))
 
     async def broadcast_to_chat(self, chat_guid: str, message: str | dict) -> None:
         if isinstance(message, dict):
@@ -100,6 +101,8 @@ class SocketManager:
                     sockets = self.chats.get(chat_guid)
                     if sockets:
                         for socket in sockets:
+                            if socket.client_state != WebSocketState.CONNECTED:
+                                continue
                             data = message['data']
                             await socket.send_text(data)
                 else:
@@ -109,7 +112,8 @@ class SocketManager:
             logger.exception(f'Exception occurred: {exc}')
 
     async def send_error(self, message: str, websocket: WebSocket):
-        await websocket.send_json({'status': 'error', 'message': message})
+        if websocket.client_state == WebSocketState.CONNECTED:
+            await websocket.send_json({'status': 'error', 'content': message})
 
     async def quit(self):
         self._task_cancel = False

core/routers/chat/routes.py

@@ -20,7 +20,7 @@ from core.db.crud import CRUD
 from core.db.engine import get_async_session, get_cache
 from core.db.models import Message, Chat, User, ReadStatus, chat_participant
 from core.routers.auth.schema import LoginUser
-from core.routers.auth.services import get_current_user
+from core.routers.auth.services import get_current_user_form_http, get_current_user_form_ws
 from core.base.exceptions import WebsocketTooManyRequests
 
 from . import schema
@@ -41,7 +41,7 @@ chat_router = APIRouter(prefix='/chat', tags=['聊天模块'])
     ]
 )
 async def get_all_chats(
-        login_user: LoginUser = Depends(get_current_user),
+        login_user: LoginUser = Depends(get_current_user_form_http),
         db_session: AsyncSession = Depends(get_async_session),
         page: int = Query(default=1, gt=0),
         page_size: int = Query(default=15, gt=0, le=200),
@@ -82,13 +82,13 @@ async def get_all_chats(
     '/{chat_guid}/messages/', summary='获取单个聊天的历史消息',
     response_model=base_schema.BaseResponse[
         base_schema.Paginated[
-            List[Dict]
+            List[schema.MessageOut]
         ]
     ]
 )
 async def get_single_chat_messages(
         chat_guid: str,
-        login_user: LoginUser = Depends(get_current_user),
+        login_user: LoginUser = Depends(get_current_user_form_http),
         db_session: AsyncSession = Depends(get_async_session),
         page: int = Query(default=1, gt=0),
         page_size: int = Query(default=15, gt=0, le=200),
@@ -125,15 +125,18 @@ async def get_single_chat_messages(
         .options(selectinload(Message.user), selectinload(Message.chat))
     )
     records = await db_session.execute(stmt)
+    results = [
+        {
+            'user': record[0].user.dict(),
+            'chat_guid': record[0].chat.guid,
+            'message_guid': record[0].guid,
+            **record[0].dict(exclude={'guid'})
+        } for record in records
+    ]
 
     return {
         'data': {
-            'list': [
-                {
-                    **card.dict(),
-                    **user_collection.dict()
-                } for card, user_collection in records
-            ],
+            'list': results,
             'total': total.scalar(),
             'page': page,
             'page_size': page_size,
@@ -144,7 +147,7 @@ async def get_single_chat_messages(
 @chat_router.post('/create', summary='创建新聊天', response_model=base_schema.BaseResponse[schema.Chat])
 async def create_new_chat(
         payload: schema.ChatCreateIn,
-        login_user: LoginUser = Depends(get_current_user),
+        login_user: LoginUser = Depends(get_current_user_form_http),
         db_session: AsyncSession = Depends(get_async_session),
 ):
     if payload.recipient_user_id == login_user.user_id:
@@ -201,7 +204,7 @@ async def create_new_chat(
 @chat_router.delete('/{chat_guid}/', summary='删除聊天', response_model=base_schema.BaseResponse)
 async def delete_single_chat(
         chat_guid: str,
-        login_user: LoginUser = Depends(get_current_user),
+        login_user: LoginUser = Depends(get_current_user_form_http),
         db_session: AsyncSession = Depends(get_async_session),
 ):
     chat = await CRUD(db_session, Chat).update(
@@ -217,17 +220,16 @@ async def delete_single_chat(
 @chat_router.websocket('/ws')
 async def get_chat_websocket(
         websocket: WebSocket,
-        login_user: LoginUser = Depends(get_current_user),
+        login_user: LoginUser = Depends(get_current_user_form_ws),
         db_session: AsyncSession = Depends(get_async_session),
-        redis_conn: aioredis.Redis = Depends(get_cache)
 ):
-    current_user = await CRUD(db_session, User).get(filters={'user_id': login_user.user_id})
     await socket_manager.connect_socket(websocket=websocket)
-    logger.info('Websocket connection is established')
-    ratelimit = WebSocketRateLimiter(times=50, seconds=10, callback=websocket_callback)
+    logger.info(f'Websocket connection is established: {login_user.user_id}')
     user_id = str(login_user.user_id)
     await socket_manager.add_user_socket_connection(user_id, websocket)
 
+    ratelimit = WebSocketRateLimiter(times=50, seconds=10, callback=websocket_callback)
+    current_user = await CRUD(db_session, User).get(filters={'user_id': login_user.user_id})
     if chats := await get_user_active_direct_chats(db_session, current_user=current_user):
         for chat_guid in chats:
             await socket_manager.add_user_to_chat(chat_guid, websocket)
@@ -238,10 +240,10 @@ async def get_chat_websocket(
         while True:
             try:
                 incoming_message = await websocket.receive_json()
+                logger.info(f'新消息：{incoming_message}')
                 await ratelimit(websocket)
-
                 try:
-                    message = schema.Message(**incoming_message)
+                    message = schema.WSMessage(**incoming_message)
                 except ValueError:
                     await socket_manager.send_error(websocket, 'The message type or format you sent is invalid')
                     continue
@@ -255,10 +257,9 @@ async def get_chat_websocket(
                 await handler(
                     websocket=websocket,
                     db_session=db_session,
-                    cache=redis_conn,
                     current_user=current_user,
                     chats=chats,
-                    message=message,
+                    incoming_message=message,
                 )
 
             except (JSONDecodeError, AttributeError) as exc:
@@ -273,4 +274,4 @@ async def get_chat_websocket(
                 logger.exception(f'User: {login_user} sent too many ws requests')
                 await websocket.send_json({'status': 'error', 'message': 'You have sent too many requests'})
     except WebSocketDisconnect:
-        logging.info('Websocket is disconnected')
+        logging.info(f'Websocket is disconnected: {login_user.user_id}')

core/routers/chat/schema.py

@@ -17,31 +17,6 @@ class ChatUser(LiteUser):
     pass
 
 
-class Message(BaseModel):
-    type: Literal[
-        'chat.new_message', 'chat.message_read'
-    ]
-    message_guid: UUID4
-    message_type: MessageType
-    chat_guid: UUID4
-    receive_user: Optional[int] = None
-    body: str
-
-
-class MessageOut(BaseModel):
-    type: Literal[
-        'chat.new_message', 'chat.send_status_reply'
-    ]
-    code: Optional[int] = 200
-    status: Optional[Literal['ok', 'error']] = 'ok'
-    message_id: int
-    message_guid: UUID4
-    message_type: MessageType
-    chat_guid: UUID4
-    body: Optional[str] = None
-    friend_user: LiteUser
-
-
 class Chat(BaseModel):
     chat_guid: UUID4
     friend_user: LiteUser
@@ -51,10 +26,51 @@ class ChatCreateIn(BaseModel):
     recipient_user_id: int
 
 
+class PreviewMessage(BaseModel):
+    message_type: MessageType
+    content: str
+
+
 class ChatUnreadOut(BaseModel):
     chat_guid: UUID4
     created_at: str
     updated_at: str
     recipient_user: LiteUser
-    last_message: str | None
+    last_message: PreviewMessage | None
     unread_messages_count: int
+
+
+class MessageOut(BaseModel):
+    chat_id: int
+    chat_guid: UUID4
+    id: int
+    message_guid: UUID4
+    message_type: MessageType
+    content: str
+    user: LiteUser
+    created_at: str
+    updated_at: str
+
+
+class WSMessage(BaseModel):
+    type: Literal[
+        'chat.new_message', 'chat.message_read', 'chat.ping'
+    ]
+    message_guid: UUID4
+    message_type: MessageType
+    chat_guid: UUID4
+    receive_user: Optional[int] = None
+    content: str
+
+
+class WSMessageOut(BaseModel):
+    type: Literal[
+        'chat.new_message', 'chat.send_status_reply',
+    ]
+    status: Optional[Literal['ok', 'error']] = 'ok'
+    message_id: int
+    message_guid: UUID4
+    message_type: MessageType
+    chat_guid: UUID4
+    content: Optional[str] = None
+    friend_user: Optional[LiteUser] = None

core/routers/chat/services.py

@@ -50,7 +50,7 @@ async def get_unread_messages_per_chat(
     # 查询消息情况
     read_status_alias = aliased(ReadStatus)
     query = (
-        select(Message.chat_id, Message.content, func.count().label("message_count"))
+        select(Message.chat_id, Message.content, Message.message_type, func.count().label("message_count"))
         .join(
             read_status_alias,
             and_(read_status_alias.user_id == login_user.user_id, read_status_alias.chat_id == Message.chat_id),
@@ -70,8 +70,11 @@ async def get_unread_messages_per_chat(
 
     for messages_info in unread_messages_info:
         unread_messages_info_per_chat[messages_info[0]] = {
-            'message': messages_info[1],
-            'message_count': messages_info[2],
+            'message': {
+                'message_type': messages_info[2],
+                'content': messages_info[1],
+            },
+            'message_count': messages_info[3],
         }
 
     results = []