Update app.py

app.py

@@ -1,3 +1,4 @@
+# app.py
 import os
 import json
 import time
@@ -8,11 +9,12 @@ import asyncio
 import aiosqlite
 import tempfile
 import io
+import logging
+from contextlib import asynccontextmanager
 
 from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect, Header, Query, Request, UploadFile, File
 from fastapi.staticfiles import StaticFiles
 from fastapi.responses import Response, JSONResponse, FileResponse
-import logging
 
 from cryptography.hazmat.primitives.ciphers.aead import AESGCM
 from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
@@ -34,85 +36,117 @@ assert len(MESSAGE_KEY) == 32
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger("InfinityChat")
 
-# --------------------------- Database ---------------------------
-async def init_db():
-    db = await aiosqlite.connect(DATABASE_URL)
-    db.row_factory = aiosqlite.Row
-    await db.execute("PRAGMA journal_mode=WAL;")
-    await db.execute("PRAGMA foreign_keys=ON;")
-    await db.execute("""
-        CREATE TABLE IF NOT EXISTS users (
-            id INTEGER PRIMARY KEY,
-            username TEXT UNIQUE NOT NULL,
-            display_name TEXT NOT NULL DEFAULT '',
-            password_hash TEXT NOT NULL,
-            salt TEXT NOT NULL,
-            avatar_path TEXT,
-            token TEXT UNIQUE,
-            created_at INTEGER DEFAULT (strftime('%s','now'))
-        )
-    """)
-    await db.execute("""
-        CREATE TABLE IF NOT EXISTS messages (
-            id INTEGER PRIMARY KEY,
-            sender_id INTEGER NOT NULL,
-            encrypted_content TEXT NOT NULL,
-            timestamp_ms INTEGER NOT NULL,
-            reply_to_id INTEGER,
-            is_edited INTEGER DEFAULT 0,
-            is_deleted INTEGER DEFAULT 0,
-            file_path TEXT,
-            file_type TEXT,
-            file_name TEXT,
-            FOREIGN KEY(sender_id) REFERENCES users(id)
-        )
-    """)
-    await db.execute("""
-        CREATE TABLE IF NOT EXISTS read_receipts (
-            user_id INTEGER NOT NULL,
-            message_id INTEGER NOT NULL,
-            PRIMARY KEY (user_id, message_id)
-        )
-    """)
-    await db.execute("CREATE INDEX IF NOT EXISTS idx_msg_time ON messages(timestamp_ms)")
-    await db.commit()
-    await db.close()
-
-asyncio.get_event_loop().run_until_complete(init_db())
+# --------------------------- Database initialization ---------------------------
+async def init_database():
+    """Initialize the SQLite database with proper async context."""
+    async with aiosqlite.connect(DATABASE_URL) as db:
+        db.row_factory = aiosqlite.Row
+        await db.execute("PRAGMA journal_mode=WAL;")
+        await db.execute("PRAGMA foreign_keys=ON;")
+        
+        await db.execute("""
+            CREATE TABLE IF NOT EXISTS users (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                username TEXT UNIQUE NOT NULL,
+                display_name TEXT NOT NULL DEFAULT '',
+                password_hash TEXT NOT NULL,
+                salt TEXT NOT NULL,
+                avatar_path TEXT,
+                token TEXT UNIQUE,
+                created_at INTEGER DEFAULT (strftime('%s','now'))
+            )
+        """)
+        await db.execute("""
+            CREATE TABLE IF NOT EXISTS messages (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                sender_id INTEGER NOT NULL,
+                encrypted_content TEXT NOT NULL,
+                timestamp_ms INTEGER NOT NULL,
+                reply_to_id INTEGER,
+                is_edited INTEGER DEFAULT 0,
+                is_deleted INTEGER DEFAULT 0,
+                file_path TEXT,
+                file_type TEXT,
+                file_name TEXT,
+                FOREIGN KEY(sender_id) REFERENCES users(id)
+            )
+        """)
+        await db.execute("""
+            CREATE TABLE IF NOT EXISTS read_receipts (
+                user_id INTEGER NOT NULL,
+                message_id INTEGER NOT NULL,
+                PRIMARY KEY (user_id, message_id)
+            )
+        """)
+        await db.execute("CREATE INDEX IF NOT EXISTS idx_msg_time ON messages(timestamp_ms)")
+        await db.commit()
+    
+    logger.info("✅ Database initialized successfully")
+
+# --------------------------- Lifespan handler (proper async startup) ---------------------------
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Handle startup and shutdown events properly."""
+    # Startup
+    await init_database()
+    yield
+    # Shutdown
+    # Cleanup if needed
+
+# --------------------------- FastAPI app with lifespan ---------------------------
+app = FastAPI(title="InfinityChat", lifespan=lifespan)
+app.mount("/static", StaticFiles(directory="static"), name="static")
 
-# --------------------------- Helpers ---------------------------
+# --------------------------- Database helpers ---------------------------
 async def get_db():
+    """Get a database connection with proper async context."""
     db = await aiosqlite.connect(DATABASE_URL)
     db.row_factory = aiosqlite.Row
     await db.execute("PRAGMA journal_mode=WAL")
+    await db.execute("PRAGMA foreign_keys=ON")
     return db
 
+# --------------------------- Encryption helpers ---------------------------
 def encrypt_message(plain: str) -> str:
+    """Encrypt message content with AES-256-GCM."""
     aesgcm = AESGCM(MESSAGE_KEY)
     nonce = os.urandom(12)
     ct = aesgcm.encrypt(nonce, plain.encode(), None)
     return base64.urlsafe_b64encode(nonce + ct).decode()
 
 def decrypt_message(encrypted_b64: str) -> str:
+    """Decrypt message content."""
     raw = base64.urlsafe_b64decode(encrypted_b64)
     nonce, ct = raw[:12], raw[12:]
     aesgcm = AESGCM(MESSAGE_KEY)
     return aesgcm.decrypt(nonce, ct, None).decode()
 
+# --------------------------- Password handling ---------------------------
 def hash_password(password: str, salt: Optional[str] = None) -> tuple[str, str]:
+    """Hash password with PBKDF2."""
     if salt is None:
         salt = os.urandom(16).hex()
-    kdf = PBKDF2HMAC(algorithm=hashes.SHA256(), length=32, salt=salt.encode(), iterations=600000)
+    kdf = PBKDF2HMAC(
+        algorithm=hashes.SHA256(),
+        length=32,
+        salt=salt.encode(),
+        iterations=600000,
+        backend=default_backend()
+    )
     return base64.urlsafe_b64encode(kdf.derive(password.encode())).decode(), salt
 
 def verify_password(password: str, salt: str, stored_hash: str) -> bool:
+    """Verify password against stored hash."""
     key, _ = hash_password(password, salt)
     return key == stored_hash
 
 def generate_token() -> str:
+    """Generate a secure random token."""
     return base64.urlsafe_b64encode(os.urandom(32)).decode()
 
+# --------------------------- User lookup ---------------------------
 async def get_user_by_token(token: str):
+    """Get user by token with proper async context."""
     db = await get_db()
     try:
         async with db.execute("SELECT * FROM users WHERE token = ?", (token,)) as cur:
@@ -120,62 +154,114 @@ async def get_user_by_token(token: str):
     finally:
         await db.close()
 
-# --------------------------- FastAPI app ---------------------------
-app = FastAPI(title="InfinityChat")
-app.mount("/static", StaticFiles(directory="static"), name="static")
+async def get_user_by_username(username: str):
+    """Get user by username."""
+    db = await get_db()
+    try:
+        async with db.execute("SELECT * FROM users WHERE username = ?", (username,)) as cur:
+            return await cur.fetchone()
+    finally:
+        await db.close()
 
 # --------------------------- HTTP endpoints ---------------------------
 @app.post("/signup")
 async def signup(username: str, password: str, display_name: str = ""):
+    """Create a new user account."""
     db = await get_db()
-    if await (await db.execute("SELECT 1 FROM users WHERE username = ?", (username,))).fetchone():
+    try:
+        existing = await get_user_by_username(username)
+        if existing:
+            raise HTTPException(400, "Username already taken")
+        
+        pwd_hash, salt = hash_password(password)
+        token = generate_token()
+        
+        await db.execute(
+            "INSERT INTO users (username, display_name, password_hash, salt, token) VALUES (?, ?, ?, ?, ?)",
+            (username, display_name, pwd_hash, salt, token)
+        )
+        await db.commit()
+        
+        return {"token": token, "username": username, "display_name": display_name}
+    finally:
         await db.close()
-        raise HTTPException(400, "Username taken")
-    pwd_hash, salt = hash_password(password)
-    token = generate_token()
-    await db.execute("INSERT INTO users (username, display_name, password_hash, salt, token) VALUES (?,?,?,?,?)",
-                     (username, display_name, pwd_hash, salt, token))
-    await db.commit()
-    await db.close()
-    return {"token": token, "username": username, "display_name": display_name}
 
 @app.post("/login")
 async def login(username: str, password: str):
+    """Authenticate user and return token."""
     db = await get_db()
-    user = await (await db.execute("SELECT * FROM users WHERE username = ?", (username,))).fetchone()
-    if not user or not verify_password(password, user["salt"], user["password_hash"]):
+    try:
+        user = await get_user_by_username(username)
+        if not user or not verify_password(password, user["salt"], user["password_hash"]):
+            raise HTTPException(400, "Invalid credentials")
+        
+        token = generate_token()
+        await db.execute("UPDATE users SET token = ? WHERE id = ?", (token, user["id"]))
+        await db.commit()
+        
+        return {
+            "token": token,
+            "username": user["username"],
+            "display_name": user["display_name"],
+            "avatar_path": user["avatar_path"]
+        }
+    finally:
         await db.close()
-        raise HTTPException(400, "Invalid credentials")
-    token = generate_token()
-    await db.execute("UPDATE users SET token = ? WHERE id = ?", (token, user["id"]))
-    await db.commit()
-    await db.close()
-    return {"token": token, "username": user["username"], "display_name": user["display_name"], "avatar_path": user["avatar_path"]}
 
 @app.get("/me")
 async def me(token: str = Header(...)):
+    """Get current user info."""
+    user = await get_user_by_token(token)
+    if not user:
+        raise HTTPException(401, "Invalid token")
+    return {
+        "id": user["id"],
+        "username": user["username"],
+        "display_name": user["display_name"],
+        "avatar_path": user["avatar_path"]
+    }
+
+@app.post("/update_profile")
+async def update_profile(display_name: str = None, token: str = Header(...)):
+    """Update user profile."""
     user = await get_user_by_token(token)
     if not user:
         raise HTTPException(401)
-    return {"id": user["id"], "username": user["username"], "display_name": user["display_name"], "avatar_path": user["avatar_path"]}
+    
+    if display_name is not None:
+        db = await get_db()
+        try:
+            await db.execute("UPDATE users SET display_name = ? WHERE id = ?", (display_name, user["id"]))
+            await db.commit()
+        finally:
+            await db.close()
+    
+    return {"status": "ok"}
 
 @app.post("/upload_avatar")
 async def upload_avatar(file: UploadFile = File(...), token: str = Header(...)):
+    """Upload user avatar."""
     user = await get_user_by_token(token)
     if not user:
         raise HTTPException(401)
+    
     data = await file.read()
     ext = os.path.splitext(file.filename)[1] if file.filename else ".jpg"
-    remote = f"avatars/{user['username']}_{uuid.uuid4().hex}{ext}"
-    store_file(remote, data)
+    remote_path = f"avatars/{user['username']}_{uuid.uuid4().hex}{ext}"
+    
+    store_file(remote_path, data)
+    
     db = await get_db()
-    await db.execute("UPDATE users SET avatar_path = ? WHERE id = ?", (remote, user["id"]))
-    await db.commit()
-    await db.close()
-    return {"avatar_path": remote}
+    try:
+        await db.execute("UPDATE users SET avatar_path = ? WHERE id = ?", (remote_path, user["id"]))
+        await db.commit()
+    finally:
+        await db.close()
+    
+    return {"avatar_path": remote_path}
 
-# Chunked upload endpoint
-upload_sessions: Dict[str, dict] = {}  # upload_id -> {chunks: dict, filename, file_type, total_chunks, user_id}
+# --------------------------- Chunked upload ---------------------------
+upload_sessions: Dict[str, dict] = {}
 
 @app.post("/upload_chunk")
 async def upload_chunk(
@@ -187,9 +273,11 @@ async def upload_chunk(
     upload_id: str = Query(...),
     token: str = Header(...)
 ):
+    """Handle chunked file upload."""
     user = await get_user_by_token(token)
     if not user:
         raise HTTPException(401)
+    
     if upload_id not in upload_sessions:
         upload_sessions[upload_id] = {
             "chunks": {},
@@ -199,65 +287,75 @@ async def upload_chunk(
             "user_id": user["id"],
             "received": 0
         }
+    
     ses = upload_sessions[upload_id]
+    
     if chunk_index in ses["chunks"]:
         return {"status": "duplicate"}
+    
     ses["chunks"][chunk_index] = await file.read()
     ses["received"] += 1
-    if ses["received"] == total_chunks:
-        # Reassemble
+    
+    if ses["received"] == ses["total_chunks"]:
+        # All chunks received, reassemble
         full_data = b"".join(ses["chunks"][i] for i in sorted(ses["chunks"]))
-        remote = f"uploads/{user['username']}_{uuid.uuid4().hex}/{file_name}"
+        remote_path = f"uploads/{user['username']}_{uuid.uuid4().hex}/{file_name}"
+        
         try:
-            # Encrypt and store via bucket
-            store_file(remote, full_data)
+            store_file(remote_path, full_data)
         except Exception as e:
             logger.error(f"Failed to store file: {e}")
             del upload_sessions[upload_id]
             raise HTTPException(500, "Storage error")
+        
         del upload_sessions[upload_id]
+        
         return {
             "status": "complete",
-            "file_path": remote,
+            "file_path": remote_path,
             "file_name": file_name,
             "file_type": file_type
         }
-    return {"status": "chunk_ok"}
+    
+    return {"status": "chunk_ok", "received": ses["received"], "total": ses["total_chunks"]}
 
-# Downloads proxy (forces auth)
+# --------------------------- File download ---------------------------
 @app.get("/download/{file_path:path}")
 async def download(file_path: str, token: str = Header(...)):
+    """Download a file (authenticated proxy)."""
     user = await get_user_by_token(token)
     if not user:
         raise HTTPException(401)
+    
     try:
         data = retrieve_file(file_path)
     except Exception as e:
         raise HTTPException(404, "File not found or decryption failed")
+    
     import mimetypes
-    mt = mimetypes.guess_type(file_path)[0] or "application/octet-stream"
-    return Response(content=data, media_type=mt)
+    media_type = mimetypes.guess_type(file_path)[0] or "application/octet-stream"
+    return Response(content=data, media_type=media_type)
 
 @app.get("/")
 async def root():
+    """Serve the main page."""
     return FileResponse("static/index.html")
 
-# --------------------------- WebSocket ---------------------------
+# --------------------------- WebSocket Manager ---------------------------
 class ConnectionManager:
     def __init__(self):
-        self.active: Dict[str, Set[WebSocket]] = {}  # username -> set of websockets
-        self.user_read: Dict[int, int] = {}          # user_id -> last message id marked read
-
+        self.active: Dict[str, Set[WebSocket]] = {}
+    
     async def connect(self, ws: WebSocket, username: str):
         await ws.accept()
         self.active.setdefault(username, set()).add(ws)
-
+    
     def disconnect(self, ws: WebSocket, username: str):
         if username in self.active:
             self.active[username].discard(ws)
             if not self.active[username]:
                 del self.active[username]
-
+    
     async def broadcast(self, msg: dict, exclude: WebSocket = None):
         for uname, ws_set in self.active.items():
             for w in list(ws_set):
@@ -267,38 +365,42 @@ class ConnectionManager:
                     await w.send_json(msg)
                 except Exception:
                     self.disconnect(w, uname)
-
+    
     def online_users(self) -> list:
         return list(self.active.keys())
 
 manager = ConnectionManager()
 
+# --------------------------- WebSocket endpoint ---------------------------
 @app.websocket("/ws")
 async def ws_endpoint(ws: WebSocket, token: str = Query(...)):
+    """WebSocket endpoint for real-time communication."""
     if not token:
         await ws.close(code=4001)
         return
+    
     user = await get_user_by_token(token)
     if not user:
         await ws.close(code=4001, reason="Invalid token")
         return
-
+    
     username = user["username"]
     uid = user["id"]
+    
     await manager.connect(ws, username)
-
-    # Send online users immediately
+    
+    # Send online users
     online = manager.online_users()
     await ws.send_json({"type": "online_users", "users": online})
     await manager.broadcast({"type": "user_online", "username": username}, exclude=ws)
-
+    
     try:
         while True:
             data = await ws.receive_json()
             mtype = data.get("type")
-
-            # ---- New message ----
+            
             if mtype == "chat_message":
+                # Handle new message
                 content = data.get("content", "")
                 reply_to = data.get("reply_to_id")
                 client_id = data.get("client_id")
@@ -306,17 +408,21 @@ async def ws_endpoint(ws: WebSocket, token: str = Query(...)):
                 file_type = data.get("file_type")
                 file_name = data.get("file_name")
                 ts = int(time.time() * 1000)
-
+                
                 encrypted = encrypt_message(content)
+                
                 db = await get_db()
-                cur = await db.execute(
-                    """INSERT INTO messages (sender_id, encrypted_content, timestamp_ms, reply_to_id, file_path, file_type, file_name)
-                       VALUES (?, ?, ?, ?, ?, ?, ?)""",
-                    (uid, encrypted, ts, reply_to, file_path, file_type, file_name))
-                mid = cur.lastrowid
-                await db.commit()
-                await db.close()
-
+                try:
+                    cur = await db.execute(
+                        """INSERT INTO messages (sender_id, encrypted_content, timestamp_ms, reply_to_id, file_path, file_type, file_name)
+                           VALUES (?, ?, ?, ?, ?, ?, ?)""",
+                        (uid, encrypted, ts, reply_to, file_path, file_type, file_name)
+                    )
+                    mid = cur.lastrowid
+                    await db.commit()
+                finally:
+                    await db.close()
+                
                 msg_out = {
                     "type": "new_message",
                     "id": mid,
@@ -335,113 +441,147 @@ async def ws_endpoint(ws: WebSocket, token: str = Query(...)):
                     "client_id": client_id,
                     "status": "sent"
                 }
-                # Deliver to sender
+                
+                # Send confirmation to sender
                 await ws.send_json({**msg_out, "status": "delivered"})
                 # Broadcast to others
                 await manager.broadcast(msg_out, exclude=ws)
-
-            # ---- Load history (cursor‑based) ----
+            
             elif mtype == "load_messages":
+                # Handle message loading with cursor-based pagination
                 before_id = data.get("before_id")
                 limit = min(data.get("limit", 50), 100)
+                
                 db = await get_db()
-                q = "SELECT * FROM messages WHERE is_deleted = 0"
-                params = []
-                if before_id is not None:
-                    q += " AND id < ?"
-                    params.append(before_id)
-                q += " ORDER BY id DESC LIMIT ?"
-                params.append(limit)
-                rows = await (await db.execute(q, tuple(params))).fetchall()
-                rows = list(reversed(rows))
-                msgs = []
-                for row in rows:
-                    try:
-                        plain = decrypt_message(row["encrypted_content"])
-                    except Exception:
-                        plain = "[decryption error]"
-                    s = await db.execute("SELECT id, username, display_name, avatar_path FROM users WHERE id = ?", (row["sender_id"],))
-                    s = await s.fetchone()
-                    msgs.append({
-                        "id": row["id"],
-                        "sender_id": row["sender_id"],
-                        "username": s["username"] if s else "unknown",
-                        "display_name": s["display_name"] if s else "",
-                        "avatar_path": s["avatar_path"] if s else "",
-                        "content": plain,
-                        "timestamp_ms": row["timestamp_ms"],
-                        "reply_to_id": row["reply_to_id"],
-                        "file_path": row["file_path"],
-                        "file_type": row["file_type"],
-                        "file_name": row["file_name"],
-                        "is_edited": bool(row["is_edited"]),
-                        "is_deleted": bool(row["is_deleted"])
-                    })
-                await db.close()
-                await ws.send_json({"type": "messages_batch", "messages": msgs, "before_id": before_id})
-
-            # ---- Edit ----
+                try:
+                    q = "SELECT * FROM messages WHERE is_deleted = 0"
+                    params = []
+                    if before_id is not None:
+                        q += " AND id < ?"
+                        params.append(before_id)
+                    q += " ORDER BY id DESC LIMIT ?"
+                    params.append(limit)
+                    
+                    rows = await (await db.execute(q, tuple(params))).fetchall()
+                    rows = list(reversed(rows))
+                    
+                    msgs = []
+                    for row in rows:
+                        try:
+                            plain = decrypt_message(row["encrypted_content"])
+                        except Exception:
+                            plain = "[decryption error]"
+                        
+                        s = await db.execute(
+                            "SELECT id, username, display_name, avatar_path FROM users WHERE id = ?",
+                            (row["sender_id"],)
+                        )
+                        s = await s.fetchone()
+                        
+                        msgs.append({
+                            "id": row["id"],
+                            "sender_id": row["sender_id"],
+                            "username": s["username"] if s else "unknown",
+                            "display_name": s["display_name"] if s else "",
+                            "avatar_path": s["avatar_path"] if s else "",
+                            "content": plain,
+                            "timestamp_ms": row["timestamp_ms"],
+                            "reply_to_id": row["reply_to_id"],
+                            "file_path": row["file_path"],
+                            "file_type": row["file_type"],
+                            "file_name": row["file_name"],
+                            "is_edited": bool(row["is_edited"]),
+                            "is_deleted": bool(row["is_deleted"])
+                        })
+                    
+                    await ws.send_json({"type": "messages_batch", "messages": msgs, "before_id": before_id})
+                finally:
+                    await db.close()
+            
             elif mtype == "edit_message":
+                # Handle message edit
                 mid, new_text = data.get("message_id"), data.get("content")
                 db = await get_db()
-                row = await db.execute("SELECT * FROM messages WHERE id = ? AND sender_id = ? AND is_deleted = 0", (mid, uid))
-                if not await row.fetchone():
+                try:
+                    row = await db.execute(
+                        "SELECT * FROM messages WHERE id = ? AND sender_id = ? AND is_deleted = 0",
+                        (mid, uid)
+                    )
+                    if await row.fetchone():
+                        new_enc = encrypt_message(new_text)
+                        await db.execute(
+                            "UPDATE messages SET encrypted_content = ?, is_edited = 1 WHERE id = ?",
+                            (new_enc, mid)
+                        )
+                        await db.commit()
+                        await manager.broadcast({"type": "edit_message", "message_id": mid, "content": new_text})
+                finally:
                     await db.close()
-                    continue
-                new_enc = encrypt_message(new_text)
-                await db.execute("UPDATE messages SET encrypted_content = ?, is_edited = 1 WHERE id = ?", (new_enc, mid))
-                await db.commit()
-                await db.close()
-                await manager.broadcast({"type": "edit_message", "message_id": mid, "content": new_text})
-
-            # ---- Delete (permanent) ----
+            
             elif mtype == "delete_message":
+                # Handle permanent message deletion
                 mid = data.get("message_id")
                 db = await get_db()
-                row = await db.execute("SELECT * FROM messages WHERE id = ? AND sender_id = ?", (mid, uid))
-                msg = await row.fetchone()
-                if not msg:
+                try:
+                    row = await db.execute("SELECT * FROM messages WHERE id = ? AND sender_id = ?", (mid, uid))
+                    msg = await row.fetchone()
+                    if msg:
+                        if msg["file_path"]:
+                            try:
+                                delete_file(msg["file_path"])
+                            except Exception:
+                                pass
+                        
+                        await db.execute("DELETE FROM messages WHERE id = ?", (mid,))
+                        await db.execute("DELETE FROM read_receipts WHERE message_id = ?", (mid,))
+                        await db.commit()
+                        await manager.broadcast({"type": "delete_message", "message_id": mid})
+                finally:
                     await db.close()
-                    continue
-                if msg["file_path"]:
-                    try:
-                        delete_file(msg["file_path"])
-                    except Exception:
-                        pass
-                await db.execute("DELETE FROM messages WHERE id = ?", (mid,))
-                await db.execute("DELETE FROM read_receipts WHERE message_id = ?", (mid,))
-                await db.commit()
-                await db.close()
-                await manager.broadcast({"type": "delete_message", "message_id": mid})
-
-            # ---- Typing ----
+            
             elif mtype == "typing":
-                await manager.broadcast({"type": "typing_update", "username": username, "is_typing": data.get("is_typing", False)}, exclude=ws)
-
-            # ---- Read receipts ----
+                # Handle typing indicator
+                await manager.broadcast(
+                    {"type": "typing_update", "username": username, "is_typing": data.get("is_typing", False)},
+                    exclude=ws
+                )
+            
             elif mtype == "mark_read":
+                # Handle read receipts
                 up_to = data.get("up_to_message_id")
-                if not up_to:
-                    continue
-                db = await get_db()
-                # Which messages sent by others, <= up_to, not deleted, and not yet marked read by this user?
-                rows = await db.execute(
-                    """SELECT id, sender_id FROM messages
-                       WHERE id <= ? AND sender_id != ? AND is_deleted = 0
-                       EXCEPT SELECT message_id FROM read_receipts WHERE user_id = ?""",
-                    (up_to, uid, uid))
-                receipts = []
-                async for r in rows:
-                    await db.execute("INSERT OR IGNORE INTO read_receipts (user_id, message_id) VALUES (?, ?)", (uid, r["id"]))
-                    receipts.append(r["id"])
-                await db.commit()
-                await db.close()
-                # Notify the original senders
-                for mid in receipts:
-                    await manager.broadcast({"type": "read_receipt", "message_id": mid, "reader_username": username})
-
+                if up_to:
+                    db = await get_db()
+                    try:
+                        rows = await db.execute(
+                            """SELECT id, sender_id FROM messages
+                               WHERE id <= ? AND sender_id != ? AND is_deleted = 0
+                               EXCEPT SELECT message_id FROM read_receipts WHERE user_id = ?""",
+                            (up_to, uid, uid)
+                        )
+                        
+                        receipts = []
+                        async for r in rows:
+                            await db.execute(
+                                "INSERT OR IGNORE INTO read_receipts (user_id, message_id) VALUES (?, ?)",
+                                (uid, r["id"])
+                            )
+                            receipts.append(r["id"])
+                        
+                        await db.commit()
+                        
+                        for mid in receipts:
+                            await manager.broadcast({
+                                "type": "read_receipt",
+                                "message_id": mid,
+                                "reader_username": username
+                            })
+                    finally:
+                        await db.close()
+    
     except WebSocketDisconnect:
         pass
+    except Exception as e:
+        logger.error(f"WebSocket error for {username}: {e}")
     finally:
         manager.disconnect(ws, username)
         if username not in manager.active: