From b44005fa5247d0355a25e89a06f6da5a63b6c605 Mon Sep 17 00:00:00 2001
From: roo hutton <roo@roohutton.com>
Date: Wed, 17 Apr 2024 08:05:17 +0100
Subject: [PATCH] Merge pull request #17879 from overleaf/rh-refresh-collab

Refresh websocket on collaborator change

GitOrigin-RevId: b05444d592f5952a08bad51c9f529ed9183d042f
---
 .../real-time/app/js/WebsocketLoadBalancer.js |  9 ++++-
 .../unit/js/WebsocketLoadBalancerTests.js     | 33 +++++++++++++++++++
 .../Collaborators/CollaboratorsController.js  |  4 +--
 .../ide-react/hooks/use-socket-listeners.ts   | 16 +++++++++
 4 files changed, 59 insertions(+), 3 deletions(-)

diff --git a/services/real-time/app/js/WebsocketLoadBalancer.js b/services/real-time/app/js/WebsocketLoadBalancer.js
index 018c2e82a6..4c9bb96b5a 100644
--- a/services/real-time/app/js/WebsocketLoadBalancer.js
+++ b/services/real-time/app/js/WebsocketLoadBalancer.js
@@ -44,6 +44,9 @@ module.exports = WebsocketLoadBalancer = {
       ) {
         return true
       }
+    } else if (message?.message === 'project:collaboratorAccessLevel:changed') {
+      const changedUserId = message.payload[0].userId
+      return userId === changedUserId
     }
     return false
   },
@@ -181,7 +184,11 @@ module.exports = WebsocketLoadBalancer = {
                 },
                 'disconnecting client'
               )
-              client.emit('project:access:revoked')
+              if (
+                message?.message !== 'project:collaboratorAccessLevel:changed'
+              ) {
+                client.emit('project:access:revoked')
+              }
               client.disconnect()
             } else {
               if (isRestrictedMessage && client.ol_context.is_restricted_user) {
diff --git a/services/real-time/test/unit/js/WebsocketLoadBalancerTests.js b/services/real-time/test/unit/js/WebsocketLoadBalancerTests.js
index 3aadbfcc12..85bd9a5cd9 100644
--- a/services/real-time/test/unit/js/WebsocketLoadBalancerTests.js
+++ b/services/real-time/test/unit/js/WebsocketLoadBalancerTests.js
@@ -69,6 +69,39 @@ describe('WebsocketLoadBalancer', function () {
       ).to.equal(false)
     })
 
+    describe('collaborator access level changed', function () {
+      const messageName = 'project:collaboratorAccessLevel:changed'
+      const client = {
+        ol_context: { user_id: 'abcd' },
+      }
+      it('should return true if the user id matches', function () {
+        const message = {
+          message: messageName,
+          payload: [
+            {
+              userId: 'abcd',
+            },
+          ],
+        }
+        expect(
+          this.WebsocketLoadBalancer.shouldDisconnectClient(client, message)
+        ).to.equal(true)
+      })
+      it('should return false if the user id does not match', function () {
+        const message = {
+          message: messageName,
+          payload: [
+            {
+              userId: 'xyz',
+            },
+          ],
+        }
+        expect(
+          this.WebsocketLoadBalancer.shouldDisconnectClient(client, message)
+        ).to.equal(false)
+      })
+    })
+
     describe('user removed from project', function () {
       const messageName = 'userRemovedFromProject'
       const client = {
diff --git a/services/web/app/src/Features/Collaborators/CollaboratorsController.js b/services/web/app/src/Features/Collaborators/CollaboratorsController.js
index d2959ecd28..d21da72751 100644
--- a/services/web/app/src/Features/Collaborators/CollaboratorsController.js
+++ b/services/web/app/src/Features/Collaborators/CollaboratorsController.js
@@ -82,8 +82,8 @@ async function setCollaboratorInfo(req, res, next) {
     )
     EditorRealTimeController.emitToRoom(
       projectId,
-      'project:membership:changed',
-      { members: true }
+      'project:collaboratorAccessLevel:changed',
+      { userId }
     )
     res.sendStatus(204)
   } catch (err) {
diff --git a/services/web/frontend/js/features/ide-react/hooks/use-socket-listeners.ts b/services/web/frontend/js/features/ide-react/hooks/use-socket-listeners.ts
index 71207f75fa..0dfd5bab86 100644
--- a/services/web/frontend/js/features/ide-react/hooks/use-socket-listeners.ts
+++ b/services/web/frontend/js/features/ide-react/hooks/use-socket-listeners.ts
@@ -58,6 +58,22 @@ function useSocketListeners() {
     )
   )
 
+  useSocketListener(
+    socket,
+    'project:collaboratorAccessLevel:changed',
+    useCallback(() => {
+      listProjectMembers(projectId)
+        .then(({ members }) => {
+          if (members) {
+            setProjectMembers(members)
+          }
+        })
+        .catch(err => {
+          debugConsole.error('Error fetching members for project', err)
+        })
+    }, [projectId, setProjectMembers])
+  )
+
   useSocketListener(
     socket,
     'project:membership:changed',