From 0cdca391497e766d2a0538d2fc1c2a3e486b32b9 Mon Sep 17 00:00:00 2001 From: Carlos Herrero <26092748+hbcarlos@users.noreply.github.com> Date: Fri, 20 Oct 2023 19:42:41 +0200 Subject: [PATCH] Improve document sessions (#210) --- jupyter_collaboration/handlers.py | 12 +++++++++++- jupyter_collaboration/rooms/base.py | 5 ++--- jupyter_collaboration/rooms/document.py | 2 ++ jupyter_collaboration/utils.py | 1 + packages/docprovider/src/utils.ts | 3 ++- packages/docprovider/src/yprovider.ts | 25 +++++++++++++++++++++---- 6 files changed, 39 insertions(+), 9 deletions(-) diff --git a/jupyter_collaboration/handlers.py b/jupyter_collaboration/handlers.py index 438a4ee7..7a00c3ff 100644 --- a/jupyter_collaboration/handlers.py +++ b/jupyter_collaboration/handlers.py @@ -22,6 +22,7 @@ JUPYTER_COLLABORATION_EVENTS_URI, LogLevel, MessageType, + RoomMessages, decode_file_path, ) @@ -150,7 +151,7 @@ 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}." ) @@ -158,6 +159,15 @@ async def open(self, room_id): 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)) diff --git a/jupyter_collaboration/rooms/base.py b/jupyter_collaboration/rooms/base.py index 3b7bd3ba..12a0699b 100644 --- a/jupyter_collaboration/rooms/base.py +++ b/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. diff --git a/jupyter_collaboration/rooms/document.py b/jupyter_collaboration/rooms/document.py index 4330ce86..e6e85ee3 100644 --- a/jupyter_collaboration/rooms/document.py +++ b/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) diff --git a/jupyter_collaboration/utils.py b/jupyter_collaboration/utils.py index cf1c7bc8..c7bfbd9b 100644 --- a/jupyter_collaboration/utils.py +++ b/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): diff --git a/packages/docprovider/src/utils.ts b/packages/docprovider/src/utils.ts index 47815952..9321995a 100644 --- a/packages/docprovider/src/utils.ts +++ b/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 } diff --git a/packages/docprovider/src/yprovider.ts b/packages/docprovider/src/yprovider.ts index c87e4484..018ee90a 100644 --- a/packages/docprovider/src/yprovider.ts +++ b/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 { - 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(); private _serverUrl: string; private _sharedModel: YDocument;