Add websocket endpoint

backend/requirements.txt

@@ -2,4 +2,5 @@ PyJWT>=2.8.0
 fastapi[standard]>=0.116.1
 pydantic>=2.11.7
 sqlalchemy>=2.0.43
-bcrypt>=4.3.0
+bcrypt>=4.3.0
+websockets>=15.0.1

backend/routes/messaging.py

@@ -1,16 +1,34 @@
 from datetime import datetime
-from mailbox import Message
-from fastapi import APIRouter, Depends, HTTPException
+from fastapi import APIRouter, Depends, HTTPException, WebSocket
 from sqlalchemy.orm import Session
 from dependencies import get_current_user, get_db
-from models import SendMessageRequest
+from models import Message, SendMessageRequest
 
 
 router = APIRouter()
 
+async def get_messages_inner(current_user: dict, db: Session):
+    messages = db.query(Message).order_by(Message.timestamp.asc()).all()
+
+    messages_data = []
+    for msg in messages:
+        messages_data.append({
+            "id": msg.id,
+            "content": msg.content,
+            "timestamp": msg.timestamp.isoformat(),
+            "is_author": msg.user_id == current_user["user_id"],
+            "is_read": msg.is_read,
+            "username": msg.author.username
+        })
+
+    return {
+        "status": "success",
+        "messages": messages_data
+    }
+
 
 @router.post("/send_message")
-def send_message(
+async def send_message(
     request: SendMessageRequest,
     current_user: dict = Depends(get_current_user),
     db: Session = Depends(get_db)
@@ -35,24 +53,29 @@ def send_message(
 
 
 @router.get("/get_messages")
-def get_messages(
+async def get_messages(
     current_user: dict = Depends(get_current_user),
     db: Session = Depends(get_db)
 ):
-    messages = db.query(Message).order_by(Message.timestamp.asc()).all()
+    return await get_messages_inner(current_user, db)
 
-    messages_data = []
-    for msg in messages:
-        messages_data.append({
-            "id": msg.id,
-            "content": msg.content,
-            "timestamp": msg.timestamp.isoformat(),
-            "is_author": msg.user_id == current_user["user_id"],
-            "is_read": msg.is_read,
-            "username": msg.author.username
-        })
 
-    return {
-        "status": "success",
-        "messages": messages_data
-    }
+@router.websocket("/chat/ws")
+async def messaging(websocket: WebSocket):
+    await websocket.accept()
+
+    while True:
+        data = await websocket.receive_json()
+
+        if data.type == "ping":
+            await websocket.send_json({"type": "ping", "data": {"status": "success"}})
+        elif data.type == "getMessages":
+            try:
+                current_user = get_current_user()
+                db = next(get_db())
+
+                await websocket.send_json({"type": "getMessages", "data": await get_messages_inner(current_user, db)})
+            except HTTPException as e:
+                await websocket.send_json({"type": "getMessages", "error": {"code": e.status_code, "detail": e.detail}})
+        else:
+            await websocket.send_json({"type": data.type, "error": {"code": 400, "detail": "Invalid type"}})