adding yjsclient and yjsclientgroup and server awareness observer

* adding yjsclient and yjsclientgroup and server awareness observer

* add websocket check

---------

Co-authored-by: Jialin Zhang <[email protected]>

Author: jzhang20133

jupyter_rtc_core/rooms/yroom.py

@@ -12,32 +12,43 @@
     from typing import Literal, Tuple
 
 class YRoom:
+    """A Room to manage all client connection to one notebook file"""
+    room_id: str
+    """Room Id"""
     ydoc: pycrdt.Doc
+    """Ydoc"""
     awareness: pycrdt.Awareness
+    """Ydoc awareness object"""
     loop: asyncio.AbstractEventLoop
+    """Event loop"""
     log: Logger
+    """Log object"""
     _client_group: YjsClientGroup
+    """Client group to manage synced and desynced clients"""
     _message_queue: asyncio.Queue[Tuple[str, bytes]]
+    """A message queue per room to keep websocket messages in order"""
 
 
-    def __init__(self, log: Logger, loop: asyncio.AbstractEventLoop):
+    def __init__(self, *, room_id: str, log: Logger, loop: asyncio.AbstractEventLoop):
         # Bind instance attributes
         self.log = log
         self.loop = loop
+        self.room_id = room_id
 
         # Initialize YDoc, YAwareness, YjsClientGroup, and message queue
         self.ydoc = pycrdt.Doc()
         self.awareness = pycrdt.Awareness(ydoc=self.ydoc)
-        self._client_group = YjsClientGroup()
+        self.awareness.observe(self.send_server_awareness)
+        self._client_group = YjsClientGroup(room_id=room_id, log=self.log, loop=self.loop)
         self._message_queue = asyncio.Queue()
+        
+        # Start observer on the `ydoc` to ensure new updates are broadcast to
+        # all clients and saved to disk.
+        self.ydoc.observe(lambda event: self.write_sync_update(event.update))
 
         # Start background task that routes new messages in the message queue
         # to the appropriate handler method.
         self.loop.create_task(self._on_new_message())
-
-        # Start observer on the `ydoc` to ensure new updates are broadcast to
-        # all clients and saved to disk.
-        self.ydoc.observe(lambda event: self.write_sync_update(event.update))
 
 
     @property
@@ -124,7 +135,7 @@ def handle_sync_step1(self, client_id: str, message: bytes) -> None:
         - Sending a new SyncStep1 message immediately after.
         """
         # Mark client as desynced
-        new_client = self.clients.get(client_id, synced_only=False)
+        new_client = self.clients.get(client_id)
         self.clients.mark_desynced(client_id)
 
         # Compute SyncStep2 reply
@@ -146,7 +157,6 @@ def handle_sync_step1(self, client_id: str, message: bytes) -> None:
             # TODO: remove the assert once websocket is made required
             assert isinstance(new_client.websocket, WebSocketHandler)
             new_client.websocket.write_message(sync_step2_message)
-            self.clients.mark_synced(client_id)
         except Exception as e:
             self.log.error(
                 "An exception occurred when writing the SyncStep2 reply "
@@ -155,6 +165,8 @@ def handle_sync_step1(self, client_id: str, message: bytes) -> None:
             self.log.exception(e)
             return
 
+        self.clients.mark_synced(client_id)
+        
         # Send SyncStep1 message
         try:
             assert isinstance(new_client.websocket, WebSocketHandler)
@@ -197,9 +209,10 @@ def handle_sync_update(self, client_id: str, message: bytes) -> None:
         clients after this method is called via the `self.write_sync_update()`
         observer.
         """
-        # Ignore the message if client is desynced
+        # Remove client and kill websocket if received SyncUpdate when client is desynced
         if self._should_ignore_update(client_id, "SyncUpdate"):
-            return
+            self.log.error(f"Should not receive SyncUpdate message when double handshake is not completed for client '{client_id}' and room '{self.room_id}'")
+            self._client_group.remove(client_id)
 
         # Apply the SyncUpdate to the YDoc
         try:
@@ -235,10 +248,6 @@ def write_sync_update(self, message_payload: bytes, client_id: str | None = None
 
 
     def handle_awareness_update(self, client_id: str, message: bytes) -> None:
-        # Ignore the message if client is desynced
-        if self._should_ignore_update(client_id, "AwarenessUpdate"):
-            return
-
         # Apply the AwarenessUpdate message
         try:
             message_payload = message[1:]
@@ -263,7 +272,7 @@ def _should_ignore_update(self, client_id: str, message_type: Literal['Awareness
         more readable warnings.
         """
 
-        client = self.clients.get(client_id, synced_only=False)
+        client = self.clients.get(client_id)
         if not client.synced:
             self.log.warning(
                 f"Ignoring a {message_type} message from client "
@@ -293,7 +302,22 @@ def _broadcast_message(self, message: bytes, message_type: Literal['AwarenessUpd
                     f"to client '{client.id}:'"
                 )
                 self.log.exception(e)
+                
+    def send_server_awareness(self, type: str, changes: tuple[dict[str, Any], Any]) -> None:
+        """
+        Callback to broadcast the server awareness to clients.
 
+        Arguments:
+            type: The change type.
+            changes: The awareness changes.
+        """
+        if type != "update" or changes[1] != "local":
+            return
+        
+        updated_clients = [v for value in changes[0].values() for v in value]
+        state = self.awareness.encode_awareness_update(updated_clients)
+        message = pycrdt.create_awareness_message(state)
+        self._broadcast_message(message, "AwarenessUpdate")
 
     def stop(self) -> None:
         # TODO: requires YRoomLoader implementation

jupyter_rtc_core/websockets/clients.py

@@ -5,31 +5,40 @@
 """
 
 from __future__ import annotations
+from datetime import timedelta, timezone
+import datetime
+from logging import Logger
 from typing import TYPE_CHECKING
 from dataclasses import dataclass
 import uuid
+import asyncio
 
 if TYPE_CHECKING:
     from tornado.websocket import WebSocketHandler
-
-@dataclass
 class YjsClient:
     """Data model that represents all data associated
     with a user connecting to a YDoc through JupyterLab."""
 
-    websocket: WebSocketHandler | None = None
-    """
-    The Tornado WebSocketHandler handling the WS connection to this
-    client.
-
-    TODO: make this required
-    """
-    
-    id: str = str(uuid.uuid4())
+    websocket: WebSocketHandler | None
+    """The Tornado WebSocketHandler handling the WS connection to this client."""
+    id: str
     """UUIDv4 string that uniquely identifies this client."""
-
-    synced: bool = False
+    synced: bool
     """Indicates whether the SS1 + SS2 handshake has been completed."""
+    last_modified: datetime
+    """Indicates the last modified time when synced state is modified"""
+
+    def __init__(self, websocket):
+        self.websocket: WebSocketHandler | None = websocket
+        self.id: str = str(uuid.uuid4())
+        self.synced: bool = False
+        self.last_modified = datetime.now(timezone.utc)
+        
+    
+    @synced.setter
+    def synced(self, v: bool):
+        self.synced = v
+        self.last_modified = datetime.now(timezone.utc)
 
 class YjsClientGroup:
     """
@@ -39,42 +48,101 @@ class YjsClientGroup:
     New clients start as desynced. Consumers should call mark_synced() to mark a
     new client as synced once the SS1 + SS2 handshake is complete.
     
-    TODO: Automatically removes desynced clients if they do not become synced after
+    Automatically removes desynced clients if they do not become synced after
     a certain timeout.
     """
+    room_id: str
+    """Room Id for associated YRoom"""
     synced: dict[str, YjsClient]
+    """A dict of client_id and synced YjsClient mapping"""
     desynced: dict[str, YjsClient]
+    """A dict of client_id and desynced YjsClient mapping"""
+    log: Logger
+    """Log object"""
+    loop: asyncio.AbstractEventLoop
+    """Event loop"""
+    poll_interval_seconds: int
+    """The poll time interval used while auto removing desynced clients"""
+    desynced_timeout_seconds: int
+    """The max time period in seconds that a desynced client does not become synced before get auto removed from desynced dict"""
+    
+    def __init__(self, *, room_id: str, log: Logger, loop: asyncio.AbstractEventLoop, poll_interval_seconds: int = 60, desynced_timeout_seconds: int = 120):
+        self.room_id = room_id
+        self.synced: dict[str, YjsClient] = {}
+        self.desynced: dict[str, YjsClient] = {}
+        self.log = log
+        self.loop = loop
+        self.loop.create_task(self._clean_desynced())
+        self.poll_interval_seconds = poll_interval_seconds
+        self.desynced_timeout_seconds = desynced_timeout_seconds
 
     def add(self, websocket: WebSocketHandler) -> str:
         """Adds a pending client to the group. Returns a client ID."""
-        return ""
+        client = YjsClient(websocket)
+        self.desynced[client.id] = client
+        return client.id
 
     def mark_synced(self, client_id: str) -> None:
         """Marks a client as synced."""
-        return
+        if client := self.desynced.pop(client_id, None):
+            client.synced = True
+            self.synced[client.id] = client
 
     def mark_desynced(self, client_id: str) -> None:
         """Marks a client as desynced."""
-        return
+        if client := self.synced.pop(client_id, None):
+            client.synced = False
+            self.desynced[client.id] = client
 
     def remove(self, client_id: str) -> None:
         """Removes a client from the group."""
-        return
+        if client := self.desynced.pop(client_id, None) is None: 
+            client = self.synced.pop(client_id, None)
+        if client and client.websocket and client.websocket.ws_connection: 
+            try:
+                client.websocket.close()
+            except Exception as e:
+                self.log.exception(f"An exception occurred when remove client '{client_id}' for room '{self.room_id}': {e}")  
 
-    def get(self, client_id: str, synced_only: bool = True) -> YjsClient:
+    def get(self, client_id: str) -> YjsClient:
         """
         Gets a client from its ID.
-        Set synced_only=False to also get desynced clients.
         """
-        return YjsClient()
+        if client_id in self.desynced: 
+            client = self.desynced[client_id]
+        if client_id in self.synced:
+            client = self.synced[client_id]
+        if client.websocket and client.websocket.ws_connection:
+            return client
+        error_message = f"The client_id '{client_id}' is not found in client group in room '{self.room_id}'"
+        self.log.error(error_message)
+        raise Exception(error_message)
 
     def get_all(self, synced_only: bool = True) -> list[YjsClient]:
         """
         Returns a list of all synced clients.
         Set synced_only=False to also get desynced clients.
         """
-        return []
+        if synced_only: 
+            return list(client for client in self.synced.values() if client.websocket and client.websocket.ws_connection)
+        return list(client for client in self.desynced.values() if client.websocket and client.websocket.ws_connection)
 
     def is_empty(self) -> bool:
         """Returns whether the client group is empty."""
-        return False
+        return len(self.synced) == 0 and len(self.desynced) == 0
+    
+    async def _clean_desynced(self) -> None:
+        while True: 
+            try:
+                await asyncio.sleep(self._poll_interval_seconds)
+                for (client_id, client) in list(self.desynced.items()): 
+                    if client.last_modified <= datetime.now(timezone.utc) - timedelta(seconds=self.desynced_timeout_seconds):
+                        self.log.warning(f"Remove client '{client_id}' for room '{self.room_id}' since client does not become synced after {self.desynced_timeout_seconds} seconds.")  
+                        self.remove(client_id)
+                for (client_id, client) in list(self.synced.items()): 
+                    if client.websocket is None or client.websocket.ws_connection is None:
+                        self.log.warning(f"Remove client '{client_id}' for room '{self.room_id}' since client does not become synced after {self.desynced_timeout_seconds} seconds.")  
+                        self.remove(client_id)
+            except asyncio.CancelledError:
+                break
+