feat: websocket rooms (#515)

apps/client/src/features/editor/title-editor.tsx

@@ -10,9 +10,7 @@ import {
   pageEditorAtom,
   titleEditorAtom,
 } from "@/features/editor/atoms/editor-atoms";
-import {
-  useUpdatePageMutation,
-} from "@/features/page/queries/page-query";
+import { useUpdatePageMutation } from "@/features/page/queries/page-query";
 import { useDebouncedValue } from "@mantine/hooks";
 import { useAtom } from "jotai";
 import { treeDataAtom } from "@/features/page/tree/atoms/tree-data-atom";
@@ -39,7 +37,11 @@ export function TitleEditor({
 }: TitleEditorProps) {
   const [debouncedTitleState, setDebouncedTitleState] = useState(null);
   const [debouncedTitle] = useDebouncedValue(debouncedTitleState, 500);
-  const updatePageMutation = useUpdatePageMutation();
+  const {
+    data: updatedPageData,
+    mutate: updatePageMutation,
+    status,
+  } = useUpdatePageMutation();
   const pageEditor = useAtomValue(pageEditorAtom);
   const [, setTitleEditor] = useAtom(titleEditorAtom);
   const [treeData, setTreeData] = useAtom(treeDataAtom);
@@ -47,7 +49,6 @@ export function TitleEditor({
   const navigate = useNavigate();
   const [activePageId, setActivePageId] = useState(pageId);
 
-
   const titleEditor = useEditor({
     extensions: [
       Document.extend({
@@ -87,24 +88,29 @@ export function TitleEditor({
 
   useEffect(() => {
     if (debouncedTitle !== null && activePageId === pageId) {
-      updatePageMutation.mutate({
+      updatePageMutation({
         pageId: pageId,
         title: debouncedTitle,
       });
+    }
+  }, [debouncedTitle]);
+
+  useEffect(() => {
+    if (status === "success" && updatedPageData) {
+      const newTreeData = updateTreeNodeName(treeData, pageId, debouncedTitle);
+      setTreeData(newTreeData);
 
       setTimeout(() => {
         emit({
           operation: "updateOne",
+          spaceId: updatedPageData.spaceId,
           entity: ["pages"],
           id: pageId,
           payload: { title: debouncedTitle, slugId: slugId },
         });
       }, 50);
-
-      const newTreeData = updateTreeNodeName(treeData, pageId, debouncedTitle);
-      setTreeData(newTreeData);
     }
-  }, [debouncedTitle]);
+  }, [updatedPageData, status]);
 
   useEffect(() => {
     if (titleEditor && title !== titleEditor.getText()) {

apps/client/src/features/page/tree/components/space-tree.tsx

@@ -133,13 +133,13 @@ export default function SpaceTree({ spaceId, readOnly }: SpaceTreeProps) {
             flatTreeItems = [
               ...flatTreeItems,
               ...children.filter(
-                (child) => !flatTreeItems.some((item) => item.id === child.id),
+                (child) => !flatTreeItems.some((item) => item.id === child.id)
               ),
             ];
           };
 
           const fetchPromises = ancestors.map((ancestor) =>
-            fetchAndUpdateChildren(ancestor),
+            fetchAndUpdateChildren(ancestor)
           );
 
           // Wait for all fetch operations to complete
@@ -153,7 +153,7 @@ export default function SpaceTree({ spaceId, readOnly }: SpaceTreeProps) {
             const updatedTree = appendNodeChildren(
               data,
               rootChild.id,
-              rootChild.children,
+              rootChild.children
             );
             setData(updatedTree);
 
@@ -248,7 +248,7 @@ function Node({ node, style, dragHandle, tree }: NodeRendererProps<any>) {
       const updatedTreeData = appendNodeChildren(
         treeData,
         node.data.id,
-        childrenTree,
+        childrenTree
       );
 
       setTreeData(updatedTreeData);
@@ -279,6 +279,7 @@ function Node({ node, style, dragHandle, tree }: NodeRendererProps<any>) {
     setTimeout(() => {
       emit({
         operation: "updateOne",
+        spaceId: node.data.spaceId,
         entity: ["pages"],
         id: node.id,
         payload: { icon: emoji.native },
@@ -293,6 +294,7 @@ function Node({ node, style, dragHandle, tree }: NodeRendererProps<any>) {
     setTimeout(() => {
       emit({
         operation: "updateOne",
+        spaceId: node.data.spaceId,
         entity: ["pages"],
         id: node.id,
         payload: { icon: null },

apps/client/src/features/page/tree/hooks/use-tree-mutation.ts

@@ -75,18 +75,19 @@ export function useTreeMutation<T>(spaceId: string) {
     setTimeout(() => {
       emit({
         operation: "addTreeNode",
+        spaceId: spaceId,
         payload: {
           parentId,
           index,
-          data
-        }
+          data,
+        },
       });
     }, 50);
 
     const pageUrl = buildPageUrl(
       spaceSlug,
       createdPage.slugId,
-      createdPage.title,
+      createdPage.title
     );
     navigate(pageUrl);
     return data;
@@ -156,18 +157,16 @@ export function useTreeMutation<T>(spaceId: string) {
       // check if the previous still has children
       // if no children left, change 'hasChildren' to false, to make the page toggle arrows work properly
       const childrenCount = previousParent.children.filter(
-        (child) => child.id !== draggedNodeId,
+        (child) => child.id !== draggedNodeId
       ).length;
       if (childrenCount === 0) {
         tree.update({
           id: previousParent.id,
-          changes: { ... previousParent.data, hasChildren: false } as any,
+          changes: { ...previousParent.data, hasChildren: false } as any,
         });
       }
     }
 
-    //console.log()
-
     setData(tree.data);
 
     const payload: IMovePage = {
@@ -182,7 +181,13 @@ export function useTreeMutation<T>(spaceId: string) {
       setTimeout(() => {
         emit({
           operation: "moveTreeNode",
-          payload: { id: draggedNodeId, parentId: args.parentId, index: args.index, position: newPosition },
+          spaceId: spaceId,
+          payload: {
+            id: draggedNodeId,
+            parentId: args.parentId,
+            index: args.index,
+            position: newPosition,
+          },
         });
       }, 50);
     } catch (error) {
@@ -214,17 +219,17 @@ export function useTreeMutation<T>(spaceId: string) {
       setData(tree.data);
 
       // navigate only if the current url is same as the deleted page
-      if (pageSlug && node.data.slugId === pageSlug.split('-')[1]) {
+      if (pageSlug && node.data.slugId === pageSlug.split("-")[1]) {
         navigate(getSpaceUrl(spaceSlug));
       }
 
       setTimeout(() => {
         emit({
           operation: "deleteTreeNode",
-          payload: { node: node.data }
+          spaceId: spaceId,
+          payload: { node: node.data },
         });
       }, 50);
-
     } catch (error) {
       console.error("Failed to delete page:", error);
     }

apps/client/src/features/websocket/types/types.ts

@@ -2,12 +2,14 @@ import { SpaceTreeNode } from "@/features/page/tree/types.ts";
 
 export type InvalidateEvent = {
   operation: "invalidate";
+  spaceId: string;
   entity: Array<string>;
   id?: string;
 };
 
 export type UpdateEvent = {
   operation: "updateOne";
+  spaceId: string;
   entity: Array<string>;
   id: string;
   payload: Partial<any>;
@@ -15,6 +17,7 @@ export type UpdateEvent = {
 
 export type DeleteEvent = {
   operation: "deleteOne";
+  spaceId: string;
   entity: Array<string>;
   id: string;
   payload?: Partial<any>;
@@ -22,6 +25,7 @@ export type DeleteEvent = {
 
 export type AddTreeNodeEvent = {
   operation: "addTreeNode";
+  spaceId: string;
   payload: {
     parentId: string;
     index: number;
@@ -31,6 +35,7 @@ export type AddTreeNodeEvent = {
 
 export type MoveTreeNodeEvent = {
   operation: "moveTreeNode";
+  spaceId: string;
   payload: {
     id: string;
     parentId: string;
@@ -41,6 +46,7 @@ export type MoveTreeNodeEvent = {
 
 export type DeleteTreeNodeEvent = {
   operation: "deleteTreeNode";
+  spaceId: string;
   payload: {
     node: SpaceTreeNode
   }

apps/server/src/ws/ws.gateway.ts

@@ -9,6 +9,7 @@ import { Server, Socket } from 'socket.io';
 import { TokenService } from '../core/auth/services/token.service';
 import { JwtType } from '../core/auth/dto/jwt-payload';
 import { OnModuleDestroy } from '@nestjs/common';
+import { SpaceMemberRepo } from '@docmost/db/repos/space/space-member.repo';
 
 @WebSocketGateway({
   cors: { origin: '*' },
@@ -17,7 +18,10 @@ import { OnModuleDestroy } from '@nestjs/common';
 export class WsGateway implements OnGatewayConnection, OnModuleDestroy {
   @WebSocketServer()
   server: Server;
-  constructor(private tokenService: TokenService) {}
+  constructor(
+    private tokenService: TokenService,
+    private spaceMemberRepo: SpaceMemberRepo,
+  ) {}
 
   async handleConnection(client: Socket, ...args: any[]): Promise<void> {
     try {
@@ -27,24 +31,43 @@ export class WsGateway implements OnGatewayConnection, OnModuleDestroy {
       if (token.type !== JwtType.ACCESS) {
         client.disconnect();
       }
+
+      const userId = token.sub;
+      const workspaceId = token.workspaceId;
+
+      const userSpaceIds = await this.spaceMemberRepo.getUserSpaceIds(userId);
+
+      const workspaceRoom = `workspace-${workspaceId}`;
+      const spaceRooms = userSpaceIds.map((id) => this.getSpaceRoomName(id));
+
+      client.join([workspaceRoom, ...spaceRooms]);
     } catch (err) {
       client.disconnect();
     }
   }
 
   @SubscribeMessage('message')
-  handleMessage(client: Socket, data: string): void {
-    client.broadcast.emit('message', data);
-  }
+  handleMessage(client: Socket, data: any): void {
+    const spaceEvents = [
+      'updateOne',
+      'addTreeNode',
+      'moveTreeNode',
+      'deleteTreeNode',
+    ];
 
-  @SubscribeMessage('messageToRoom')
-  handleSendMessageToRoom(@MessageBody() message: any) {
-    this.server.to(message?.roomId).emit('messageToRoom', message);
+    if (spaceEvents.includes(data?.operation) && data?.spaceId) {
+      const room = this.getSpaceRoomName(data.spaceId);
+      client.broadcast.to(room).emit('message', data);
+      return;
+    }
+
+    client.broadcast.emit('message', data);
   }
 
   @SubscribeMessage('join-room')
   handleJoinRoom(client: Socket, @MessageBody() roomName: string): void {
-    client.join(roomName);
+    // if room is a space, check if user has permissions
+    //client.join(roomName);
   }
 
   @SubscribeMessage('leave-room')
@@ -57,4 +80,8 @@ export class WsGateway implements OnGatewayConnection, OnModuleDestroy {
       this.server.close();
     }
   }
+
+  getSpaceRoomName(spaceId: string): string {
+    return `space-${spaceId}`;
+  }
 }