Files
ds-chat/backend/app/ws/chat.py
T
ksmithandClaude Sonnet 5 68e487e5ec Add unread message indicators for rooms (#38)
Shows a dot on rooms with unread messages in the sidebar, updated live
over WebSocket. Reuses the same offline-member audience computation
already used for push notifications: a member gets the real-time signal
whenever they aren't currently connected to that room's channel, which
correctly covers both "room not open" and "room open but tab
backgrounded" (the client leaves a room's channel while hidden).

Persisted server-side via a new room_memberships.last_read_at column so
state survives reload and stays consistent across devices, advanced by
an explicit mark-read call the frontend makes on room-open and on each
live message received while the room is genuinely visible -- gated on a
live visibility check, not a cached ref, so a backgrounded-but-open room
keeps accumulating unread instead of auto-marking-read the instant a
message arrives somewhere it can't be seen.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-08-16 20:44:20 -06:00

253 lines
11 KiB
Python

import uuid
from fastapi import APIRouter, Depends, WebSocket, WebSocketDisconnect
from pydantic import BaseModel, ValidationError
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from app.database import get_db
from app.models import ApiToken, Message, MessageFile, MessageImage, RoomMembership, User
from app.services.bot_service import resolve_token
from app.services.message_events import (
broadcast_member_updated,
broadcast_message_update,
broadcast_new_message,
broadcast_reaction_update,
)
from app.services.message_service import (
MessageNotFoundError,
NotMessageAuthorError,
create_message,
edit_message,
toggle_reaction,
)
from app.services.room_service import mark_room_read
router = APIRouter(tags=["ws"])
WS_UNAUTHENTICATED = 4401
class ClientEnvelope(BaseModel):
type: str
room_id: uuid.UUID | None = None
content: str | None = None
image_id: uuid.UUID | None = None
file_id: uuid.UUID | None = None
message_id: uuid.UUID | None = None
emoji: str | None = None
async def _is_room_member(db: AsyncSession, room_id: uuid.UUID, user_id: uuid.UUID) -> bool:
result = await db.execute(
select(RoomMembership).where(
RoomMembership.room_id == room_id, RoomMembership.user_id == user_id
)
)
return result.scalar_one_or_none() is not None
def _missing_scope(api_token: ApiToken | None, scope: str) -> bool:
return api_token is not None and scope not in api_token.scopes
@router.websocket("/ws/chat")
async def chat_endpoint(websocket: WebSocket, db: AsyncSession = Depends(get_db)) -> None:
api_token: ApiToken | None = None
# Bots authenticate by setting Authorization on the WS handshake itself
# (not a browser cookie) -- same connection type/endpoint a human client
# uses, just a different credential.
auth_header = websocket.headers.get("authorization")
if auth_header and auth_header.lower().startswith("bearer "):
resolved = await resolve_token(db, auth_header[len("bearer ") :].strip())
if resolved is None:
await websocket.close(code=WS_UNAUTHENTICATED)
return
user, api_token = resolved
else:
user_id_raw = websocket.session.get("user_id")
if not user_id_raw:
await websocket.close(code=WS_UNAUTHENTICATED)
return
user = await db.get(User, uuid.UUID(user_id_raw))
if user is None or not user.is_active:
await websocket.close(code=WS_UNAUTHENTICATED)
return
await websocket.accept()
manager = websocket.app.state.connection_manager
presence = websocket.app.state.presence
global_presence = websocket.app.state.global_presence
broadcaster = websocket.app.state.broadcaster
joined_rooms: set[uuid.UUID] = set()
manager.register_user(user.id, websocket)
# Only broadcast on a genuine offline->online transition (this user's
# first open connection), not for every extra tab -- broadcast_member_
# updated tells every room this user's in to refresh, which would be
# wasted churn on a transition that didn't actually change anything
# visible.
if await global_presence.connect(user.id):
await broadcast_member_updated(db, broadcaster, user.id)
try:
while True:
raw = await websocket.receive_json()
try:
envelope = ClientEnvelope.model_validate(raw)
except ValidationError:
await websocket.send_json({"type": "error", "detail": "Malformed message"})
continue
if envelope.type == "join":
if envelope.room_id is None:
await websocket.send_json({"type": "error", "detail": "room_id required"})
continue
if not await _is_room_member(db, envelope.room_id, user.id):
await websocket.send_json(
{"type": "error", "detail": "Not a member of this room"}
)
continue
manager.join(envelope.room_id, websocket)
await presence.join(envelope.room_id, user.id)
joined_rooms.add(envelope.room_id)
await websocket.send_json({"type": "joined", "room_id": str(envelope.room_id)})
elif envelope.type == "leave":
if envelope.room_id is None:
await websocket.send_json({"type": "error", "detail": "room_id required"})
continue
manager.leave(envelope.room_id, websocket)
await presence.leave(envelope.room_id, user.id)
joined_rooms.discard(envelope.room_id)
elif envelope.type == "message":
if envelope.room_id is None or (
not envelope.content
and envelope.image_id is None
and envelope.file_id is None
):
await websocket.send_json(
{
"type": "error",
"detail": "room_id and content or image_id or file_id required",
}
)
continue
if _missing_scope(api_token, "write:messages"):
await websocket.send_json(
{"type": "error", "detail": "Token missing required scope: write:messages"}
)
continue
if envelope.room_id not in joined_rooms or not await _is_room_member(
db, envelope.room_id, user.id
):
await websocket.send_json(
{"type": "error", "detail": "Not a member of this room"}
)
continue
image_id = None
if envelope.image_id is not None:
image = await db.get(MessageImage, envelope.image_id)
if image is None or image.room_id != envelope.room_id:
await websocket.send_json({"type": "error", "detail": "Invalid image"})
continue
image_id = image.id
file_id = None
if envelope.file_id is not None:
message_file = await db.get(MessageFile, envelope.file_id)
if message_file is None or message_file.room_id != envelope.room_id:
await websocket.send_json({"type": "error", "detail": "Invalid file"})
continue
file_id = message_file.id
message = await create_message(
db, envelope.room_id, user.id, envelope.content, image_id, file_id
)
# Sending implies having seen the room as of now -- without
# this, GET /rooms/mine would show the sender's own room as
# unread the instant they send into it (last_read_at isn't
# otherwise bumped until the frontend's own message echo
# triggers a mark-read call, which is a real but avoidable
# race). Deliberately not done in create_message() itself:
# the incoming-webhook path also calls it, and a webhook's
# attributed sender may not actually be watching.
await mark_room_read(db, envelope.room_id, user.id)
await broadcast_new_message(db, broadcaster, presence, envelope.room_id, message, user)
elif envelope.type == "edit":
if envelope.room_id is None or envelope.message_id is None or not envelope.content:
await websocket.send_json(
{"type": "error", "detail": "room_id, message_id, and content required"}
)
continue
if _missing_scope(api_token, "write:messages"):
await websocket.send_json(
{"type": "error", "detail": "Token missing required scope: write:messages"}
)
continue
if envelope.room_id not in joined_rooms or not await _is_room_member(
db, envelope.room_id, user.id
):
await websocket.send_json(
{"type": "error", "detail": "Not a member of this room"}
)
continue
try:
message = await edit_message(db, envelope.message_id, user.id, envelope.content)
except MessageNotFoundError:
await websocket.send_json({"type": "error", "detail": "Message not found"})
continue
except NotMessageAuthorError:
await websocket.send_json(
{"type": "error", "detail": "You can only edit your own messages"}
)
continue
await broadcast_message_update(db, broadcaster, envelope.room_id, message)
elif envelope.type == "reaction":
if (
envelope.room_id is None
or envelope.message_id is None
or not envelope.emoji
or len(envelope.emoji) > 8
):
await websocket.send_json(
{"type": "error", "detail": "room_id, message_id, and emoji required"}
)
continue
if _missing_scope(api_token, "write:messages"):
await websocket.send_json(
{"type": "error", "detail": "Token missing required scope: write:messages"}
)
continue
if envelope.room_id not in joined_rooms or not await _is_room_member(
db, envelope.room_id, user.id
):
await websocket.send_json(
{"type": "error", "detail": "Not a member of this room"}
)
continue
target_message = await db.get(Message, envelope.message_id)
if target_message is None or target_message.room_id != envelope.room_id:
await websocket.send_json({"type": "error", "detail": "Message not found"})
continue
reactions = await toggle_reaction(db, envelope.message_id, user.id, envelope.emoji)
await broadcast_reaction_update(
broadcaster, envelope.room_id, envelope.message_id, reactions
)
else:
await websocket.send_json(
{"type": "error", "detail": f"Unknown message type: {envelope.type}"}
)
except WebSocketDisconnect:
pass
finally:
manager.leave_all(websocket)
manager.unregister_user(user.id, websocket)
for room_id in joined_rooms:
await presence.leave(room_id, user.id)
if await global_presence.disconnect(user.id):
await broadcast_member_updated(db, broadcaster, user.id)