fix(collab): fix disconnect message not showing

Author: kokerinks

services/collaboration-service/app/core/connection_manager.py

@@ -2,7 +2,7 @@
 from fastapi import WebSocket, WebSocketDisconnect
 import uuid 
 import asyncio
-from app.schemas.messages import ConnectMessage, DisconnectMessage, DisplayMessage
+from app.schemas.messages import CollaboratorConnectMessage, CollaboratorDisconnectMessage, DisplayMessage
 
 class ConnectionManager:
     def __init__(self):
@@ -40,7 +40,11 @@ async def on_connect(self, session_id: str, user_id: str, out_q: asyncio.Queue):
         if collaborator_q is None:
             return None
 
-        await collaborator_q.put(ConnectMessage())
+        ## Notify collaborator that this user has connected
+        await collaborator_q.put(CollaboratorConnectMessage())
+        
+        ## Notify this user that the other user has already connected
+        await out_q.put(CollaboratorConnectMessage())
 
     async def on_disconnect(self, session_id: str, user_id: str):
         if session_id not in self.active_connections:
@@ -56,7 +60,7 @@ async def on_disconnect(self, session_id: str, user_id: str):
         if collaborator_q is None:
             return None
 
-        await collaborator_q.put(DisconnectMessage())
+        await collaborator_q.put(CollaboratorDisconnectMessage())
 
     async def on_message(self, session_id: str, user_id: str, msg: str):
         collaborator_q = self._get_collaborator_q(session_id, user_id)

services/collaboration-service/app/schemas/messages.py

@@ -4,10 +4,10 @@
 class Message(BaseModel):
     pass 
 
-class ConnectMessage(Message):
+class CollaboratorConnectMessage(Message):
     pass
 
-class DisconnectMessage(Message):
+class CollaboratorDisconnectMessage(Message):
     pass
 
 class DisplayMessage(Message):

services/collaboration-service/app/ws/sessions.py

@@ -2,7 +2,7 @@
 from fastapi import APIRouter, Query, WebSocket, WebSocketDisconnect
 from app.core.connection_manager import connection_manager
 from app.core.users import UserState
-from app.schemas.messages import Message, ConnectMessage, DisconnectMessage, DisplayMessage
+from app.schemas.messages import Message, CollaboratorConnectMessage, CollaboratorDisconnectMessage, DisplayMessage
 
 from datetime import datetime
 import asyncio
@@ -36,19 +36,22 @@ async def receiver():
                 msg = await ws.receive_text()
                 await connection_manager.on_message(session_id, user_id, msg)
         except WebSocketDisconnect:
-            print("aandle disconnect")
-            pass ##
+            try:
+                print("handle disconnect")
+                await connection_manager.on_disconnect(session_id, user_id)
+            except ValueError as e:
+                print(f"Error handling disconnect: {e}")
 
     receiver_task = asyncio.create_task(receiver())
 
     try:
         while True:
             msg = await in_q.get()
 
-            if isinstance(msg, ConnectMessage):
+            if isinstance(msg, CollaboratorConnectMessage):
                 user_status = UserState.AWAIT_POLLING
                 await ws.send_text(f"Connected to session {session_id} as user {user_id}. You can start collaborating!")
-            elif isinstance(msg, DisconnectMessage):
+            elif isinstance(msg, CollaboratorDisconnectMessage):
                 user_status = UserState.AWAIT_CONNECT
                 await ws.send_text(f"Collaborator disconnected, awaiting reconnection...")
             elif isinstance(msg, DisplayMessage):
@@ -58,12 +61,12 @@ async def receiver():
 
             await asyncio.sleep(0.1)
 
-    except WebSocketDisconnect as e:
-        try:
-            print("handle disconnect")
-            await connection_manager.on_disconnect(session_id, user_id)
-        except ValueError as e:
-            print(f"Error handling disconnect: {e}")
+    # except WebSocketDisconnect as e:
+    #     try:
+    #         print("handle disconnect")
+    #         await connection_manager.on_disconnect(session_id, user_id)
+    #     except ValueError as e:
+    #         print(f"Error handling disconnect: {e}")
     finally:
         receiver_task.cancel()