342 lines
14 KiB
Python
342 lines
14 KiB
Python
import asyncio
|
||
from typing import Annotated, overload
|
||
|
||
from app.database.chat import ChannelType, ChatChannel, ChatChannelResp, ChatMessageResp
|
||
from app.database.notification import UserNotification, insert_notification
|
||
from app.database.user import User
|
||
from app.dependencies.database import (
|
||
DBFactory,
|
||
Redis,
|
||
get_db_factory,
|
||
get_redis,
|
||
with_db,
|
||
)
|
||
from app.dependencies.user import get_current_user_and_token
|
||
from app.log import log
|
||
from app.models.chat import ChatEvent
|
||
from app.models.notification import NotificationDetail
|
||
from app.service.subscribers.chat import ChatSubscriber
|
||
from app.utils import bg_tasks
|
||
|
||
from fastapi import APIRouter, Depends, Header, Query, WebSocket, WebSocketDisconnect
|
||
from fastapi.security import SecurityScopes
|
||
from fastapi.websockets import WebSocketState
|
||
from sqlmodel import select
|
||
from sqlmodel.ext.asyncio.session import AsyncSession
|
||
|
||
logger = log("NotificationServer")
|
||
|
||
|
||
class ChatServer:
|
||
def __init__(self):
|
||
self.connect_client: dict[int, WebSocket] = {}
|
||
self.channels: dict[int, list[int]] = {}
|
||
self.redis: Redis = get_redis()
|
||
|
||
self.tasks: set[asyncio.Task] = set()
|
||
self.ChatSubscriber = ChatSubscriber()
|
||
self.ChatSubscriber.chat_server = self
|
||
self._subscribed = False
|
||
|
||
def connect(self, user_id: int, client: WebSocket):
|
||
self.connect_client[user_id] = client
|
||
|
||
def get_user_joined_channel(self, user_id: int) -> list[int]:
|
||
return [channel_id for channel_id, users in self.channels.items() if user_id in users]
|
||
|
||
async def disconnect(self, user: User, session: AsyncSession):
|
||
user_id = user.id
|
||
if user_id in self.connect_client:
|
||
del self.connect_client[user_id]
|
||
|
||
# 创建频道ID列表的副本以避免在迭代过程中修改字典
|
||
channel_ids_to_process = []
|
||
for channel_id, channel in self.channels.items():
|
||
if user_id in channel:
|
||
channel_ids_to_process.append(channel_id)
|
||
|
||
# 现在安全地处理每个频道
|
||
for channel_id in channel_ids_to_process:
|
||
# 再次检查用户是否仍在频道中（防止并发修改）
|
||
if channel_id in self.channels and user_id in self.channels[channel_id]:
|
||
self.channels[channel_id].remove(user_id)
|
||
# 使用明确的查询避免延迟加载
|
||
db_channel = (
|
||
await session.exec(select(ChatChannel).where(ChatChannel.channel_id == channel_id))
|
||
).first()
|
||
if db_channel:
|
||
await self.leave_channel(user, db_channel, session)
|
||
|
||
@overload
|
||
async def send_event(self, client: int, event: ChatEvent): ...
|
||
|
||
@overload
|
||
async def send_event(self, client: WebSocket, event: ChatEvent): ...
|
||
|
||
async def send_event(self, client: WebSocket | int, event: ChatEvent):
|
||
if isinstance(client, int):
|
||
client_ = self.connect_client.get(client)
|
||
if client_ is None:
|
||
return
|
||
client = client_
|
||
if client.client_state == WebSocketState.CONNECTED:
|
||
await client.send_text(event.model_dump_json())
|
||
|
||
async def broadcast(self, channel_id: int, event: ChatEvent):
|
||
users_in_channel = self.channels.get(channel_id, [])
|
||
logger.info(f"Broadcasting to channel {channel_id}, users: {users_in_channel}")
|
||
|
||
# 如果频道中没有用户，检查是否是多人游戏频道
|
||
if not users_in_channel:
|
||
try:
|
||
async with with_db() as session:
|
||
channel = await session.get(ChatChannel, channel_id)
|
||
if channel and channel.type == ChannelType.MULTIPLAYER:
|
||
logger.warning(
|
||
f"No users in multiplayer channel {channel_id}, message will not be delivered to anyone"
|
||
)
|
||
# 对于多人游戏房间，这可能是正常的（用户都离开了房间）
|
||
# 但我们仍然记录这个情况以便调试
|
||
except Exception as e:
|
||
logger.error(f"Failed to check channel type for {channel_id}: {e}")
|
||
|
||
for user_id in users_in_channel:
|
||
await self.send_event(user_id, event)
|
||
logger.debug(f"Sent event to user {user_id} in channel {channel_id}")
|
||
|
||
async def mark_as_read(self, channel_id: int, user_id: int, message_id: int):
|
||
await self.redis.set(f"chat:{channel_id}:last_read:{user_id}", message_id)
|
||
|
||
async def send_message_to_channel(self, message: ChatMessageResp, is_bot_command: bool = False):
|
||
logger.info(
|
||
f"Sending message to channel {message.channel_id}, message_id: "
|
||
f"{message.message_id}, is_bot_command: {is_bot_command}"
|
||
)
|
||
|
||
event = ChatEvent(
|
||
event="chat.message.new",
|
||
data={"messages": [message], "users": [message.sender]},
|
||
)
|
||
if is_bot_command:
|
||
logger.info(f"Sending bot command to user {message.sender_id}")
|
||
bg_tasks.add_task(self.send_event, message.sender_id, event)
|
||
else:
|
||
# 总是广播消息，无论是临时ID还是真实ID
|
||
logger.info(f"Broadcasting message to all users in channel {message.channel_id}")
|
||
bg_tasks.add_task(
|
||
self.broadcast,
|
||
message.channel_id,
|
||
event,
|
||
)
|
||
|
||
# 只有真实消息 ID（正数且非零）才进行标记已读和设置最后消息
|
||
# Redis 消息系统生成的ID都是正数，所以这里应该都能正常处理
|
||
if message.message_id and message.message_id > 0:
|
||
await self.mark_as_read(message.channel_id, message.sender_id, message.message_id)
|
||
await self.redis.set(f"chat:{message.channel_id}:last_msg", message.message_id)
|
||
logger.info(f"Updated last message ID for channel {message.channel_id} to {message.message_id}")
|
||
else:
|
||
logger.debug(f"Skipping last message update for message ID: {message.message_id}")
|
||
|
||
async def batch_join_channel(self, users: list[User], channel: ChatChannel, session: AsyncSession):
|
||
channel_id = channel.channel_id
|
||
|
||
not_joined = []
|
||
|
||
if channel_id not in self.channels:
|
||
self.channels[channel_id] = []
|
||
for user in users:
|
||
if user.id not in self.channels[channel_id]:
|
||
self.channels[channel_id].append(user.id)
|
||
not_joined.append(user)
|
||
|
||
for user in not_joined:
|
||
channel_resp = await ChatChannelResp.from_db(
|
||
channel,
|
||
session,
|
||
user,
|
||
self.redis,
|
||
self.channels[channel_id] if channel.type != ChannelType.PUBLIC else None,
|
||
)
|
||
await self.send_event(
|
||
user.id,
|
||
ChatEvent(
|
||
event="chat.channel.join",
|
||
data=channel_resp.model_dump(),
|
||
),
|
||
)
|
||
|
||
async def join_channel(self, user: User, channel: ChatChannel, session: AsyncSession) -> ChatChannelResp:
|
||
user_id = user.id
|
||
channel_id = channel.channel_id
|
||
|
||
if channel_id not in self.channels:
|
||
self.channels[channel_id] = []
|
||
if user_id not in self.channels[channel_id]:
|
||
self.channels[channel_id].append(user_id)
|
||
|
||
channel_resp = await ChatChannelResp.from_db(
|
||
channel,
|
||
session,
|
||
user,
|
||
self.redis,
|
||
self.channels[channel_id] if channel.type != ChannelType.PUBLIC else None,
|
||
)
|
||
|
||
await self.send_event(
|
||
user_id,
|
||
ChatEvent(
|
||
event="chat.channel.join",
|
||
data=channel_resp.model_dump(),
|
||
),
|
||
)
|
||
|
||
return channel_resp
|
||
|
||
async def leave_channel(self, user: User, channel: ChatChannel, session: AsyncSession) -> None:
|
||
user_id = user.id
|
||
channel_id = channel.channel_id
|
||
|
||
if channel_id in self.channels and user_id in self.channels[channel_id]:
|
||
self.channels[channel_id].remove(user_id)
|
||
|
||
if (c := self.channels.get(channel_id)) is not None and not c:
|
||
del self.channels[channel_id]
|
||
|
||
channel_resp = await ChatChannelResp.from_db(
|
||
channel,
|
||
session,
|
||
user,
|
||
self.redis,
|
||
self.channels.get(channel_id) if channel.type != ChannelType.PUBLIC else None,
|
||
)
|
||
await self.send_event(
|
||
user_id,
|
||
ChatEvent(
|
||
event="chat.channel.part",
|
||
data=channel_resp.model_dump(),
|
||
),
|
||
)
|
||
|
||
async def join_room_channel(self, channel_id: int, user_id: int):
|
||
async with with_db() as session:
|
||
# 使用明确的查询避免延迟加载
|
||
db_channel = (await session.exec(select(ChatChannel).where(ChatChannel.channel_id == channel_id))).first()
|
||
if db_channel is None:
|
||
logger.warning(f"Attempted to join non-existent channel {channel_id} by user {user_id}")
|
||
return
|
||
|
||
user = await session.get(User, user_id)
|
||
if user is None:
|
||
logger.warning(f"Attempted to join channel {channel_id} by non-existent user {user_id}")
|
||
return
|
||
|
||
logger.info(f"User {user_id} joining channel {channel_id} (type: {db_channel.type.value})")
|
||
await self.join_channel(user, db_channel, session)
|
||
|
||
async def leave_room_channel(self, channel_id: int, user_id: int):
|
||
async with with_db() as session:
|
||
# 使用明确的查询避免延迟加载
|
||
db_channel = (await session.exec(select(ChatChannel).where(ChatChannel.channel_id == channel_id))).first()
|
||
if db_channel is None:
|
||
logger.warning(f"Attempted to leave non-existent channel {channel_id} by user {user_id}")
|
||
return
|
||
|
||
user = await session.get(User, user_id)
|
||
if user is None:
|
||
logger.warning(f"Attempted to leave channel {channel_id} by non-existent user {user_id}")
|
||
return
|
||
|
||
logger.info(f"User {user_id} leaving channel {channel_id} (type: {db_channel.type.value})")
|
||
await self.leave_channel(user, db_channel, session)
|
||
|
||
async def new_private_notification(self, detail: NotificationDetail):
|
||
async with with_db() as session:
|
||
id = await insert_notification(session, detail)
|
||
users = (await session.exec(select(UserNotification).where(UserNotification.notification_id == id))).all()
|
||
for user_notification in users:
|
||
data = user_notification.notification.model_dump()
|
||
data["is_read"] = user_notification.is_read
|
||
data["details"] = user_notification.notification.details
|
||
await server.send_event(
|
||
user_notification.user_id,
|
||
ChatEvent(
|
||
event="new",
|
||
data=data,
|
||
),
|
||
)
|
||
|
||
|
||
server = ChatServer()
|
||
|
||
chat_router = APIRouter(include_in_schema=False)
|
||
|
||
|
||
async def _listen_stop(ws: WebSocket, user_id: int, factory: DBFactory):
|
||
try:
|
||
while True:
|
||
packets = await ws.receive_json()
|
||
if packets.get("event") == "chat.end":
|
||
async for session in factory():
|
||
user = await session.get(User, user_id)
|
||
if user is None:
|
||
break
|
||
await server.disconnect(user, session)
|
||
await ws.close(code=1000)
|
||
break
|
||
except WebSocketDisconnect as e:
|
||
logger.info(f"Client {user_id} disconnected: {e.code}, {e.reason}")
|
||
except RuntimeError as e:
|
||
if "disconnect message" in str(e):
|
||
logger.info(f"Client {user_id} closed the connection.")
|
||
else:
|
||
logger.exception(f"RuntimeError in client {user_id}: {e}")
|
||
except Exception:
|
||
logger.exception(f"Error in client {user_id}")
|
||
|
||
|
||
@chat_router.websocket("/notification-server")
|
||
async def chat_websocket(
|
||
websocket: WebSocket,
|
||
factory: Annotated[DBFactory, Depends(get_db_factory)],
|
||
token: Annotated[str | None, Query(description="认证令牌，支持通过URL参数传递")] = None,
|
||
access_token: Annotated[str | None, Query(description="访问令牌，支持通过URL参数传递")] = None,
|
||
authorization: Annotated[str | None, Header(description="Bearer认证头")] = None,
|
||
):
|
||
if not server._subscribed:
|
||
server._subscribed = True
|
||
await server.ChatSubscriber.start_subscribe()
|
||
|
||
async for session in factory():
|
||
# 优先使用查询参数中的token，支持token或access_token参数名
|
||
auth_token = token or access_token
|
||
if not auth_token and authorization:
|
||
auth_token = authorization.removeprefix("Bearer ")
|
||
|
||
if not auth_token:
|
||
await websocket.close(code=1008, reason="Missing authentication token")
|
||
return
|
||
|
||
if (
|
||
user_and_token := await get_current_user_and_token(
|
||
session, SecurityScopes(scopes=["chat.read"]), token_pw=auth_token
|
||
)
|
||
) is None:
|
||
await websocket.close(code=1008, reason="Invalid or expired token")
|
||
return
|
||
|
||
await websocket.accept()
|
||
login = await websocket.receive_json()
|
||
if login.get("event") != "chat.start":
|
||
await websocket.close(code=1008)
|
||
return
|
||
user = user_and_token[0]
|
||
user_id = user.id
|
||
server.connect(user_id, websocket)
|
||
# 使用明确的查询避免延迟加载
|
||
db_channel = (await session.exec(select(ChatChannel).where(ChatChannel.channel_id == 1))).first()
|
||
if db_channel is not None:
|
||
await server.join_channel(user, db_channel, session)
|
||
|
||
await _listen_stop(websocket, user_id, factory)
|