feat(chat): support public channel chat

app/router/chat/__init__.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+from app.config import settings
+from app.router.v2 import api_v2_router as router
+
+from . import channel, message  # noqa: F401
+from .server import chat_router as chat_router
+
+from fastapi import Query
+
+__all__ = ["chat_router"]
+
+
+@router.get("/notifications")
+async def get_notifications(max_id: int | None = Query(None)):
+    if settings.server_url is not None:
+        notification_endpoint = f"{settings.server_url}notification-server".replace(
+            "http://", "ws://"
+        ).replace("https://", "wss://")
+    else:
+        notification_endpoint = "/notification-server"
+
+    return {
+        "has_more": False,
+        "notifications": [],
+        "unread_count": 0,
+        "notification_endpoint": notification_endpoint,
+    }

app/router/chat/channel.py

@@ -0,0 +1,138 @@
+from __future__ import annotations
+
+from typing import Any
+
+from app.database.chat import (
+    ChannelType,
+    ChatChannel,
+    ChatChannelResp,
+)
+from app.database.lazer_user import User, UserResp
+from app.dependencies.database import get_db, get_redis
+from app.dependencies.user import get_current_user
+from app.router.v2 import api_v2_router as router
+
+from .server import server
+
+from fastapi import Depends, HTTPException, Query, Security
+from pydantic import BaseModel, Field
+from redis.asyncio import Redis
+from sqlmodel import select
+from sqlmodel.ext.asyncio.session import AsyncSession
+
+
+class UpdateResponse(BaseModel):
+    presence: list[ChatChannelResp] = Field(default_factory=list)
+    silences: list[Any] = Field(default_factory=list)
+
+
+@router.get("/chat/updates", response_model=UpdateResponse)
+async def get_update(
+    history_since: int | None = Query(None),
+    since: int | None = Query(None),
+    current_user: User = Security(get_current_user, scopes=["chat.read"]),
+    session: AsyncSession = Depends(get_db),
+    includes: list[str] = Query(["presence"], alias="includes[]"),
+    redis: Redis = Depends(get_redis),
+):
+    resp = UpdateResponse()
+    if "presence" in includes:
+        channel_ids = server.get_user_joined_channel(current_user.id)
+        for channel_id in channel_ids:
+            channel = await ChatChannel.get(channel_id, session)
+            if channel:
+                resp.presence.append(
+                    await ChatChannelResp.from_db(
+                        channel,
+                        session,
+                        server.channels.get(channel_id, []),
+                        current_user,
+                        redis,
+                    )
+                )
+    return resp
+
+
+@router.put("/chat/channels/{channel}/users/{user}", response_model=ChatChannelResp)
+async def join_channel(
+    channel: str,
+    user: str,
+    current_user: User = Security(get_current_user, scopes=["chat.write_manage"]),
+    session: AsyncSession = Depends(get_db),
+):
+    db_channel = await ChatChannel.get(channel, session)
+
+    if db_channel is None:
+        raise HTTPException(status_code=404, detail="Channel not found")
+    return await server.join_channel(current_user, db_channel, session)
+
+
+@router.delete(
+    "/chat/channels/{channel}/users/{user}",
+    status_code=204,
+)
+async def leave_channel(
+    channel: str,
+    user: str,
+    current_user: User = Security(get_current_user, scopes=["chat.write_manage"]),
+    session: AsyncSession = Depends(get_db),
+):
+    db_channel = await ChatChannel.get(channel, session)
+
+    if db_channel is None:
+        raise HTTPException(status_code=404, detail="Channel not found")
+    await server.leave_channel(current_user, db_channel, session)
+    return
+
+
+@router.get("/chat/channels")
+async def get_channel_list(
+    current_user: User = Security(get_current_user, scopes=["chat.read"]),
+    session: AsyncSession = Depends(get_db),
+    redis: Redis = Depends(get_redis),
+):
+    channels = (
+        await session.exec(
+            select(ChatChannel).where(ChatChannel.type == ChannelType.PUBLIC)
+        )
+    ).all()
+    results = []
+    for channel in channels:
+        assert channel.channel_id is not None
+        results.append(
+            await ChatChannelResp.from_db(
+                channel,
+                session,
+                server.channels.get(channel.channel_id, []),
+                current_user,
+                redis,
+            )
+        )
+    return results
+
+
+class GetChannelResp(BaseModel):
+    channel: ChatChannelResp
+    users: list[UserResp] = Field(default_factory=list)
+
+
+@router.get("/chat/channels/{channel}")
+async def get_channel(
+    channel: str,
+    current_user: User = Security(get_current_user, scopes=["chat.read"]),
+    session: AsyncSession = Depends(get_db),
+    redis: Redis = Depends(get_redis),
+):
+    db_channel = await ChatChannel.get(channel, session)
+    if db_channel is None:
+        raise HTTPException(status_code=404, detail="Channel not found")
+    assert db_channel.channel_id is not None
+    return GetChannelResp(
+        channel=await ChatChannelResp.from_db(
+            db_channel,
+            session,
+            server.channels.get(db_channel.channel_id, []),
+            current_user,
+            redis,
+        )
+    )

app/router/chat/message.py

@@ -0,0 +1,98 @@
+from __future__ import annotations
+
+from app.database import ChatMessageResp
+from app.database.chat import ChatChannel, ChatMessage, MessageType
+from app.database.lazer_user import User
+from app.dependencies.database import get_db
+from app.dependencies.param import BodyOrForm
+from app.dependencies.user import get_current_user
+from app.router.v2 import api_v2_router as router
+
+from .server import server
+
+from fastapi import Depends, HTTPException, Query, Security
+from pydantic import BaseModel
+from sqlmodel import col, select
+from sqlmodel.ext.asyncio.session import AsyncSession
+
+
+@router.post("/chat/ack")
+async def keep_alive(
+    history_since: int | None = Query(None),
+    since: int | None = Query(None),
+    current_user: User = Security(get_current_user, scopes=["chat.read"]),
+    session: AsyncSession = Depends(get_db),
+):
+    return {"silences": []}
+
+
+class MessageReq(BaseModel):
+    message: str
+    is_action: bool = False
+    uuid: str | None = None
+
+
+@router.post("/chat/channels/{channel}/messages", response_model=ChatMessageResp)
+async def send_message(
+    channel: str,
+    req: MessageReq = Depends(BodyOrForm(MessageReq)),
+    current_user: User = Security(get_current_user, scopes=["chat.write"]),
+    session: AsyncSession = Depends(get_db),
+):
+    db_channel = await ChatChannel.get(channel, session)
+    if db_channel is None:
+        raise HTTPException(status_code=404, detail="Channel not found")
+    msg = ChatMessage(
+        channel_id=db_channel.channel_id,
+        content=req.message,
+        sender_id=current_user.id,
+        type=MessageType.ACTION if req.is_action else MessageType.PLAIN,
+        uuid=req.uuid,
+    )
+    session.add(msg)
+    await session.commit()
+    await session.refresh(msg)
+    await session.refresh(current_user)
+    resp = await ChatMessageResp.from_db(msg, session, current_user)
+    await server.send_message_to_channel(resp)
+    return resp
+
+
+@router.get("/chat/channels/{channel}/messages", response_model=list[ChatMessageResp])
+async def get_message(
+    channel: str,
+    limit: int = Query(50, ge=1, le=50),
+    since: int = Query(default=0, ge=0),
+    until: int | None = Query(None),
+    current_user: User = Security(get_current_user, scopes=["chat.read"]),
+    session: AsyncSession = Depends(get_db),
+):
+    db_channel = await ChatChannel.get(channel, session)
+    if db_channel is None:
+        raise HTTPException(status_code=404, detail="Channel not found")
+    messages = await session.exec(
+        select(ChatMessage)
+        .where(
+            ChatMessage.channel_id == db_channel.channel_id,
+            col(ChatMessage.message_id) > since,
+            col(ChatMessage.message_id) < until if until is not None else True,
+        )
+        .order_by(col(ChatMessage.timestamp).desc())
+        .limit(limit)
+    )
+    resp = [await ChatMessageResp.from_db(msg, session) for msg in messages]
+    resp.reverse()
+    return resp
+
+
+@router.put("/chat/channels/{channel}/mark-as-read/{message}", status_code=204)
+async def mark_as_read(
+    channel: str,
+    message: int,
+    current_user: User = Security(get_current_user, scopes=["chat.read"]),
+    session: AsyncSession = Depends(get_db),
+):
+    db_channel = await ChatChannel.get(channel, session)
+    if db_channel is None:
+        raise HTTPException(status_code=404, detail="Channel not found")
+    await server.mark_as_read(db_channel.channel_id, message)

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)