Improve document sessions (#210)

jupyter_collaboration/handlers.py

@@ -22,6 +22,7 @@
     JUPYTER_COLLABORATION_EVENTS_URI,
     LogLevel,
     MessageType,
+    RoomMessages,
     decode_file_path,
 )
 
@@ -150,14 +151,23 @@ async def open(self, room_id):
 
         # Close the connection if the document session expired
         session_id = self.get_query_argument("sessionId", None)
-        if session_id and session_id != self.room.session_id:
+        if session_id is not None and session_id != self.room.session_id:
             self.log.error(
                 f"Client tried to connect to {self._room_id} with an expired session ID {session_id}."
             )
             self.close(
                 4002,
                 f"Document session {session_id} expired. You need to reload this browser tab.",
             )
+        elif session_id is None and self.room.session_id is not None:
+            # If session_id is None is because is a new document
+            # send the new session token
+            data = self.room.session_id.encode("utf8")
+            await self.send(
+                bytes([MessageType.ROOM, RoomMessages.SESSION_TOKEN])
+                + write_var_uint(len(data))
+                + data
+            )
 
         # Start processing messages in the room
         self._serve_task = asyncio.create_task(self.room.serve(self))

jupyter_collaboration/rooms/base.py

@@ -4,7 +4,6 @@
 from __future__ import annotations
 
 import asyncio
-import uuid
 from logging import Logger
 
 from ..stores import BaseYStore
@@ -15,7 +14,7 @@ class BaseRoom(YRoom):
     def __init__(self, room_id: str, store: BaseYStore | None = None, log: Logger | None = None):
         super().__init__(ready=False, ystore=store, log=log)
         self._room_id = room_id
-        self._session_id: str = str(uuid.uuid4())
+        self._session_id: str | None = None
 
     @property
     def room_id(self) -> str:
@@ -25,7 +24,7 @@ def room_id(self) -> str:
         return self._room_id
 
     @property
-    def session_id(self) -> str:
+    def session_id(self) -> str | None:
         """
         A unique identifier for the updates.
 

jupyter_collaboration/rooms/document.py

@@ -44,6 +44,7 @@ def __init__(
 
         self._file_format: str = file_format
         self._file_type: str = file_type
+        self._session_id = str(uuid.uuid4())
         self._last_modified: Any = None
         self._file: FileLoader = file
         self._document = YDOCS.get(self._file_type, YFILE)(self.ydoc)
@@ -128,6 +129,7 @@ async def initialize(self) -> None:
                     self._document.source = model["content"]
 
                     if self.ystore is not None:
+                        assert self.session_id
                         await self.ystore.create(self._room_id, self.session_id)
                         await self.ystore.encode_state_as_update(self._room_id, self.ydoc)
 

jupyter_collaboration/utils.py

@@ -22,6 +22,7 @@ class RoomMessages(IntEnum):
     FILE_CHANGED = 2
     FILE_OVERWRITTEN = 3
     DOC_OVERWRITTEN = 4
+    SESSION_TOKEN = 5
 
 
 class LogLevel(Enum):

packages/docprovider/src/utils.ts

@@ -13,5 +13,6 @@ export enum RoomMessage {
   OVERWRITE = 1,
   FILE_CHANGED = 2,
   FILE_OVERWRITTEN = 3,
-  DOC_OVERWRITTEN = 4
+  DOC_OVERWRITTEN = 4,
+  SESSION_TOKEN = 5
 }

packages/docprovider/src/yprovider.ts

@@ -15,10 +15,11 @@ import { DocumentChange, YDocument } from '@jupyter/ydoc';
 
 import * as decoding from 'lib0/decoding';
 import * as encoding from 'lib0/encoding';
+import * as url from 'lib0/url';
 import { Awareness } from 'y-protocols/awareness';
 import { WebsocketProvider as YWebsocketProvider } from 'y-websocket';
 
-import { requestDocSession } from './requests';
+import { ISessionModel, requestDocSession } from './requests';
 import { MessageType, RoomMessage } from './utils';
 
 /**
@@ -46,6 +47,7 @@ export class WebSocketProvider implements IDocumentProvider {
   constructor(options: WebSocketProvider.IOptions) {
     this._isDisposed = false;
     this._path = options.path;
+    this._session = null;
     this._contentType = options.contentType;
     this._format = options.format;
     this._serverUrl = options.url;
@@ -95,18 +97,20 @@ export class WebSocketProvider implements IDocumentProvider {
   }
 
   private async _connect(): Promise<void> {
-    const session = await requestDocSession(
+    this._session = await requestDocSession(
       this._format,
       this._contentType,
       this._path
     );
 
     const params =
-      session.sessionId !== null ? { sessionId: session.sessionId } : undefined;
+      this._session.sessionId !== null
+        ? { sessionId: this._session.sessionId }
+        : undefined;
 
     this._yWebsocketProvider = new YWebsocketProvider(
       this._serverUrl,
-      `${session.format}:${session.type}:${session.fileId}`,
+      `${this._session.format}:${this._session.type}:${this._session.fileId}`,
       this._sharedModel.ydoc,
       {
         disableBc: true,
@@ -169,6 +173,9 @@ export class WebSocketProvider implements IDocumentProvider {
           this._dialog = null;
         }
         break;
+      case RoomMessage.SESSION_TOKEN:
+        this._handleSessionToken(data);
+        break;
     }
   }
 
@@ -192,6 +199,15 @@ export class WebSocketProvider implements IDocumentProvider {
     });
   }
 
+  private _handleSessionToken(data: string): void {
+    if (this._yWebsocketProvider && this._session) {
+      const room = `${this._session.format}:${this._session.type}:${this._session.fileId}`;
+      const encodedParams = url.encodeQueryParams({ sessionId: data });
+      this._yWebsocketProvider.url =
+        this._serverUrl + '/' + room + '?' + encodedParams;
+    }
+  }
+
   private _sendReloadMsg(data: string): void {
     const encoder = encoding.createEncoder();
     encoding.writeVarUint(encoder, MessageType.ROOM);
@@ -214,6 +230,7 @@ export class WebSocketProvider implements IDocumentProvider {
   private _format: string;
   private _isDisposed: boolean;
   private _path: string;
+  private _session: ISessionModel | null;
   private _ready = new PromiseDelegate<void>();
   private _serverUrl: string;
   private _sharedModel: YDocument<DocumentChange>;