Files
ds-chat/backend/app/ws/chat.py
T
ksmithandClaude Sonnet 5 7dcc7104df Push a live signal when a user is added to a room (#26)
Previously GET /api/rooms/mine was only ever fetched once at app mount,
so a room added mid-session stayed invisible until a full page reload
-- add_member had no way to reach an already-open client at all.

Backend: ConnectionManager and Broadcaster (renamed from RoomBroadcaster)
now support per-user channels alongside the existing per-room ones, so a
signal can reach a user's socket even for a room they haven't joined
(and by definition can't have, until this fires). add_member publishes
a room_added event on the target user's channel.

Frontend: the WebSocket connection is no longer scoped to whichever
room is open -- ChatShellPage now owns one persistent connection for
the whole session (including while no room is open, which is exactly
when this bug showed), and ChatPane joins/leaves rooms on top of it.
A room_added event triggers a room-list refetch with no reload needed.

Verified end-to-end in the browser: a user sitting on the empty room
list saw a newly-added room appear live, then chatted in it normally.

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

232 lines
9.7 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_message_update,
broadcast_new_message,
broadcast_reaction_update,
)
from app.services.message_service import (
MessageNotFoundError,
NotMessageAuthorError,
create_message,
edit_message,
toggle_reaction,
)
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
broadcaster = websocket.app.state.broadcaster
joined_rooms: set[uuid.UUID] = set()
manager.register_user(user.id, websocket)
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
)
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)