feat(chat): support public channel chat

app/router/chat/server.py

@@ -0,0 +1,190 @@
+from __future__ import annotations
+
+import asyncio
+
+from app.database.chat import ChatChannel, ChatChannelResp, ChatMessageResp
+from app.database.lazer_user import User
+from app.dependencies.database import DBFactory, get_db_factory, get_redis
+from app.dependencies.user import get_current_user
+from app.log import logger
+from app.models.chat import ChatEvent
+
+from fastapi import APIRouter, Depends, Header, WebSocket, WebSocketDisconnect
+from fastapi.security import SecurityScopes
+from fastapi.websockets import WebSocketState
+from redis.asyncio import Redis
+from sqlmodel.ext.asyncio.session import AsyncSession
+
+
+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()
+
+    def _add_task(self, task):
+        task = asyncio.create_task(task)
+        self.tasks.add(task)
+        task.add_done_callback(self.tasks.discard)
+
+    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]
+        for channel_id, channel in self.channels.items():
+            if user_id in channel:
+                channel.remove(user_id)
+                channel = await ChatChannel.get(channel_id, session)
+                if channel:
+                    await self.leave_channel(user, channel, session)
+
+    async def send_event(self, client: WebSocket, event: ChatEvent):
+        if client.client_state == WebSocketState.CONNECTED:
+            await client.send_text(event.model_dump_json())
+
+    async def broadcast(self, channel_id: int, event: ChatEvent):
+        for user_id in self.channels.get(channel_id, []):
+            client = self.connect_client.get(user_id)
+            if client:
+                await self.send_event(client, event)
+
+    async def mark_as_read(self, channel_id: int, message_id: int):
+        await self.redis.set(f"chat:{channel_id}:last_msg", message_id)
+
+    async def send_message_to_channel(self, message: ChatMessageResp):
+        self._add_task(
+            self.broadcast(
+                message.channel_id,
+                ChatEvent(
+                    event="chat.message.new",
+                    data={"messages": [message], "users": [message.sender]},
+                ),
+            )
+        )
+        await self.mark_as_read(message.channel_id, message.message_id)
+
+    async def join_channel(
+        self, user: User, channel: ChatChannel, session: AsyncSession
+    ) -> ChatChannelResp:
+        user_id = user.id
+        channel_id = channel.channel_id
+        assert channel_id is not None
+
+        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, self.channels[channel_id], user, self.redis
+        )
+
+        client = self.connect_client.get(user_id)
+        if client:
+            await self.send_event(
+                client,
+                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
+        assert channel_id is not None
+
+        if channel_id in self.channels and user_id in self.channels[channel_id]:
+            self.channels[channel_id].remove(user_id)
+
+        if not self.channels.get(channel_id):
+            del self.channels[channel_id]
+
+        channel_resp = await ChatChannelResp.from_db(
+            channel, session, self.channels.get(channel_id, []), user, self.redis
+        )
+        client = self.connect_client.get(user_id)
+        if client:
+            await self.send_event(
+                client,
+                ChatEvent(
+                    event="chat.channel.part",
+                    data=channel_resp.model_dump(),
+                ),
+            )
+
+
+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"[NotificationServer] Client {user_id} disconnected: {e.code}, {e.reason}"
+        )
+    except RuntimeError as e:
+        if "disconnect message" in str(e):
+            logger.info(f"[NotificationServer] 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,
+    authorization: str = Header(...),
+    factory: DBFactory = Depends(get_db_factory),
+):
+    async for session in factory():
+        token = authorization[7:]
+        if (
+            user := await get_current_user(
+                SecurityScopes(scopes=["chat.read"]), session, token_pw=token
+            )
+        ) is None:
+            await websocket.close(code=1008)
+            return
+
+        await websocket.accept()
+        login = await websocket.receive_json()
+        if login.get("event") != "chat.start":
+            await websocket.close(code=1008)
+            return
+        user_id = user.id
+        assert user_id
+        server.connect(user_id, websocket)
+        channel = await ChatChannel.get(1, session)
+        if channel is not None:
+            await server.join_channel(user, channel, session)
+    await _listen_stop(websocket, user_id, factory)