Merge pull request #12 from docmost/websocket

Websocket sync

apps/client/src/App.tsx

@@ -13,8 +13,46 @@ import Groups from "@/pages/settings/group/groups";
 import GroupInfo from "./pages/settings/group/group-info";
 import Spaces from "@/pages/settings/space/spaces.tsx";
 import { Error404 } from "@/components/ui/error-404.tsx";
+import { useQuerySubscription } from "@/features/websocket/use-query-subscription.ts";
+import { useAtom, useAtomValue } from "jotai";
+import { socketAtom } from "@/features/websocket/atoms/socket-atom.ts";
+import { useTreeSocket } from "@/features/websocket/use-tree-socket.ts";
+import { useEffect } from "react";
+import { io } from "socket.io-client";
+import { authTokensAtom } from "@/features/auth/atoms/auth-tokens-atom.ts";
+import { SOCKET_URL } from "@/features/websocket/types";
 
 export default function App() {
+  const [, setSocket] = useAtom(socketAtom);
+  const authToken = useAtomValue(authTokensAtom);
+
+  useEffect(() => {
+    if (!authToken?.accessToken) {
+      return;
+    }
+    const newSocket = io(SOCKET_URL, {
+      transports: ["websocket"],
+      auth: {
+        token: authToken.accessToken,
+      },
+    });
+
+    // @ts-ignore
+    setSocket(newSocket);
+
+    newSocket.on("connect", () => {
+      console.log("ws connected");
+    });
+
+    return () => {
+      console.log("ws disconnected");
+      newSocket.disconnect();
+    };
+  }, [authToken?.accessToken]);
+
+  useQuerySubscription();
+  useTreeSocket();
+
   return (
     <>
       <Routes>

apps/client/src/components/layouts/components/breadcrumb.tsx

@@ -17,7 +17,7 @@ import { SpaceTreeNode } from "@/features/page/tree/types.ts";
 
 function getTitle(name: string, icon: string) {
   if (icon) {
-    return `${icon}      ${name}`;
+    return `${icon} ${name}`;
   }
   return name;
 }

apps/client/src/features/editor/extensions/extensions.ts

@@ -17,6 +17,11 @@ import { CollaborationCursor } from "@tiptap/extension-collaboration-cursor";
 import { HocuspocusProvider } from "@hocuspocus/provider";
 import { Comment, TrailingNode } from "@docmost/editor-ext";
 import GlobalDragHandle from "tiptap-extension-global-drag-handle";
+import {
+  randomElement,
+  userColors,
+} from "@/features/editor/extensions/utils.ts";
+import { IUser } from "@/features/user/types/user.types.ts";
 
 export const mainExtensions = [
   StarterKit.configure({
@@ -54,13 +59,17 @@ export const mainExtensions = [
   }),
 ] as any;
 
-type CollabExtensions = (provider: HocuspocusProvider) => any[];
+type CollabExtensions = (provider: HocuspocusProvider, user: IUser) => any[];
 
-export const collabExtensions: CollabExtensions = (provider) => [
+export const collabExtensions: CollabExtensions = (provider, user) => [
   Collaboration.configure({
     document: provider.document,
   }),
   CollaborationCursor.configure({
     provider,
+    user: {
+      name: user.name,
+      color: randomElement(userColors),
+    },
   }),
 ];

apps/client/src/features/editor/extensions/utils.ts

@@ -0,0 +1,13 @@
+export const userColors = [
+  "#958DF1",
+  "#F98181",
+  "#FBBC88",
+  "#FAF594",
+  "#70CFF8",
+  "#94FADB",
+  "#B9F18D",
+];
+
+export function randomElement(array: Array<any>) {
+  return array[Math.floor(Math.random() * array.length)];
+}

apps/client/src/features/editor/hooks/use-collaboration-url.ts

@@ -13,7 +13,9 @@ const useCollaborationURL = (): string => {
   }
   */
 
-  const API_URL = window.location.protocol + "//" + window.location.host;
+  const API_URL = import.meta.env.DEV
+    ? "http://localhost:3000"
+    : window.location.protocol + "//" + window.location.host;
 
   const wsProtocol = API_URL.startsWith("https") ? "wss" : "ws";
   return `${wsProtocol}://${API_URL.split("://")[1]}${PATH}`;

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

@@ -22,19 +22,6 @@ import CommentDialog from "@/features/comment/components/comment-dialog";
 import EditorSkeleton from "@/features/editor/components/editor-skeleton";
 import { EditorBubbleMenu } from "@/features/editor/components/bubble-menu/bubble-menu";
 
-const colors = [
-  "#958DF1",
-  "#F98181",
-  "#FBBC88",
-  "#FAF594",
-  "#70CFF8",
-  "#94FADB",
-  "#B9F18D",
-];
-const getRandomElement = (list) =>
-  list[Math.floor(Math.random() * list.length)];
-const getRandomColor = () => getRandomElement(colors);
-
 interface PageEditorProps {
   pageId: string;
   editable?: boolean;
@@ -95,7 +82,10 @@ export default function PageEditor({
     };
   }, [remoteProvider, localProvider]);
 
-  const extensions = [...mainExtensions, ...collabExtensions(remoteProvider)];
+  const extensions = [
+    ...mainExtensions,
+    ...collabExtensions(remoteProvider, currentUser.user),
+  ];
 
   const editor = useEditor(
     {
@@ -123,16 +113,6 @@ export default function PageEditor({
     [pageId, editable, remoteProvider],
   );
 
-  useEffect(() => {
-    if (editor && currentUser.user) {
-      editor
-        .chain()
-        .focus()
-        .updateUser({ ...currentUser.user, color: getRandomColor() })
-        .run();
-    }
-  }, [editor, currentUser.user]);
-
   const handleActiveCommentEvent = (event) => {
     const { commentId } = event.detail;
     setActiveCommentId(commentId);

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

@@ -15,6 +15,8 @@ import { useDebouncedValue } from "@mantine/hooks";
 import { useAtom } from "jotai";
 import { treeDataAtom } from "@/features/page/tree/atoms/tree-data-atom";
 import { updateTreeNodeName } from "@/features/page/tree/utils";
+import { useQueryEmit } from "@/features/websocket/use-query-emit.ts";
+import { History } from "@tiptap/extension-history";
 
 export interface TitleEditorProps {
   pageId: string;
@@ -28,6 +30,7 @@ export function TitleEditor({ pageId, title }: TitleEditorProps) {
   const pageEditor = useAtomValue(pageEditorAtom);
   const [, setTitleEditor] = useAtom(titleEditorAtom);
   const [treeData, setTreeData] = useAtom(treeDataAtom);
+  const emit = useQueryEmit();
 
   const titleEditor = useEditor({
     extensions: [
@@ -41,6 +44,9 @@ export function TitleEditor({ pageId, title }: TitleEditorProps) {
       Placeholder.configure({
         placeholder: "Untitled",
       }),
+      History.configure({
+        depth: 20,
+      }),
     ],
     onCreate({ editor }) {
       if (editor) {
@@ -59,6 +65,15 @@ export function TitleEditor({ pageId, title }: TitleEditorProps) {
     if (debouncedTitle !== "") {
       updatePageMutation.mutate({ pageId, title: debouncedTitle });
 
+      setTimeout(() => {
+        emit({
+          operation: "updateOne",
+          entity: ["pages"],
+          id: pageId,
+          payload: { title: debouncedTitle },
+        });
+      }, 50);
+
       const newTreeData = updateTreeNodeName(treeData, pageId, debouncedTitle);
       setTreeData(newTreeData);
     }

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

@@ -41,6 +41,7 @@ import { queryClient } from "@/main.tsx";
 import { OpenMap } from "react-arborist/dist/main/state/open-slice";
 import { useElementSize, useMergedRef } from "@mantine/hooks";
 import { dfs } from "react-arborist/dist/module/utils";
+import { useQueryEmit } from "@/features/websocket/use-query-emit.ts";
 
 interface SpaceTreeProps {
   spaceId: string;
@@ -205,6 +206,7 @@ function Node({ node, style, dragHandle, tree }: NodeRendererProps<any>) {
   const navigate = useNavigate();
   const updatePageMutation = useUpdatePageMutation();
   const [treeData, setTreeData] = useAtom(treeDataAtom);
+  const emit = useQueryEmit();
 
   async function handleLoadChildren(node: NodeApi<SpaceTreeNode>) {
     if (!node.data.hasChildren) return;
@@ -255,11 +257,29 @@ function Node({ node, style, dragHandle, tree }: NodeRendererProps<any>) {
   const handleEmojiSelect = (emoji: { native: string }) => {
     handleUpdateNodeIcon(node.id, emoji.native);
     updatePageMutation.mutateAsync({ pageId: node.id, icon: emoji.native });
+
+    setTimeout(() => {
+      emit({
+        operation: "updateOne",
+        entity: ["pages"],
+        id: node.id,
+        payload: { icon: emoji.native },
+      });
+    }, 50);
   };
 
   const handleRemoveEmoji = () => {
     handleUpdateNodeIcon(node.id, null);
     updatePageMutation.mutateAsync({ pageId: node.id, icon: null });
+
+    setTimeout(() => {
+      emit({
+        operation: "updateOne",
+        entity: ["pages"],
+        id: node.id,
+        payload: { icon: null },
+      });
+    }, 50);
   };
 
   if (node.willReceiveDrop && node.isClosed) {

apps/client/src/features/websocket/atoms/socket-atom.ts

@@ -0,0 +1,4 @@
+import { atom } from "jotai";
+import { Socket } from "socket.io-client";
+
+export const socketAtom = atom<Socket | null>(null);

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

@@ -0,0 +1,3 @@
+export const SOCKET_URL = import.meta.env.DEV
+  ? "http://localhost:3000"
+  : undefined;

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

@@ -0,0 +1,2 @@
+export * from "./types.ts";
+export * from "./constants.ts";

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

@@ -0,0 +1,14 @@
+export type InvalidateEvent = {
+  operation: "invalidate";
+  entity: Array<string>;
+  id?: string;
+};
+
+export type UpdateEvent = {
+  operation: "updateOne";
+  entity: Array<string>;
+  id: string;
+  payload: Partial<any>;
+};
+
+export type WebSocketEvent = InvalidateEvent | UpdateEvent;

apps/client/src/features/websocket/use-query-emit.ts

@@ -0,0 +1,11 @@
+import { socketAtom } from "@/features/websocket/atoms/socket-atom.ts";
+import { useAtom } from "jotai";
+import { WebSocketEvent } from "@/features/websocket/types";
+
+export const useQueryEmit = () => {
+  const [socket] = useAtom(socketAtom);
+
+  return (input: WebSocketEvent) => {
+    socket?.emit("message", input);
+  };
+};

apps/client/src/features/websocket/use-query-subscription.ts

@@ -0,0 +1,43 @@
+import React from "react";
+import { socketAtom } from "@/features/websocket/atoms/socket-atom.ts";
+import { useAtom } from "jotai";
+import { useQueryClient } from "@tanstack/react-query";
+import { WebSocketEvent } from "@/features/websocket/types";
+
+export const useQuerySubscription = () => {
+  const queryClient = useQueryClient();
+  const [socket] = useAtom(socketAtom);
+
+  React.useEffect(() => {
+    socket?.on("message", (event) => {
+      const data: WebSocketEvent = event;
+
+      switch (data.operation) {
+        case "invalidate":
+          queryClient.invalidateQueries({
+            queryKey: [...data.entity, data.id].filter(Boolean),
+          });
+          break;
+        case "updateOne":
+          queryClient.setQueryData([...data.entity, data.id], {
+            ...queryClient.getQueryData([...data.entity, data.id]),
+            ...data.payload,
+          });
+
+          /*
+          queryClient.setQueriesData(
+            { queryKey: [data.entity, data.id] },
+            (oldData: any) => {
+              const update = (entity: Record<string, unknown>) =>
+                entity.id === data.id ? { ...entity, ...data.payload } : entity;
+              return Array.isArray(oldData)
+                ? oldData.map(update)
+                : update(oldData as Record<string, unknown>);
+            },
+          );
+      */
+          break;
+      }
+    });
+  }, [queryClient, socket]);
+};

apps/client/src/features/websocket/use-tree-socket.ts

@@ -0,0 +1,62 @@
+import { useEffect, useRef } from "react";
+import { socketAtom } from "@/features/websocket/atoms/socket-atom.ts";
+import { useAtom } from "jotai";
+import { treeDataAtom } from "@/features/page/tree/atoms/tree-data-atom.ts";
+import {
+  updateTreeNodeIcon,
+  updateTreeNodeName,
+} from "@/features/page/tree/utils";
+import { WebSocketEvent } from "@/features/websocket/types";
+import { SpaceTreeNode } from "@/features/page/tree/types.ts";
+
+export const useTreeSocket = () => {
+  const [socket] = useAtom(socketAtom);
+  const [treeData, setTreeData] = useAtom(treeDataAtom);
+
+  const initialTreeData = useRef(treeData);
+
+  useEffect(() => {
+    initialTreeData.current = treeData;
+  }, [treeData]);
+
+  useEffect(() => {
+    socket?.on("message", (event) => {
+      const data: WebSocketEvent = event;
+
+      const initialData = initialTreeData.current;
+      switch (data.operation) {
+        case "invalidate":
+          // nothing to do here
+          break;
+        case "updateOne":
+          // Get the initial value of treeData
+          if (initialData && initialData.length > 0) {
+            let newTreeData: SpaceTreeNode[];
+
+            if (data.entity[0] === "pages") {
+              if (data.payload?.title !== undefined) {
+                newTreeData = updateTreeNodeName(
+                  initialData,
+                  data.id,
+                  data.payload.title,
+                );
+              }
+
+              if (data.payload?.icon !== undefined) {
+                newTreeData = updateTreeNodeIcon(
+                  initialData,
+                  data.id,
+                  data.payload.icon,
+                );
+              }
+
+              if (newTreeData && newTreeData.length > 0) {
+                setTreeData(newTreeData);
+              }
+            }
+          }
+          break;
+      }
+    });
+  }, [socket]);
+};

apps/client/src/lib/api-client.ts

@@ -2,8 +2,9 @@ import axios, { AxiosInstance } from "axios";
 import Cookies from "js-cookie";
 import Routes from "@/lib/routes";
 
+const baseUrl = import.meta.env.DEV ? "http://localhost:3000" : "";
 const api: AxiosInstance = axios.create({
-  baseURL: "/api",
+  baseURL: baseUrl + "/api",
 });
 
 api.interceptors.request.use(

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

@@ -1,27 +1,44 @@
 import {
   OnGatewayConnection,
-  OnGatewayInit,
   SubscribeMessage,
   WebSocketGateway,
   WebSocketServer,
 } from '@nestjs/websockets';
-import { Server } from 'socket.io';
+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';
 
-@WebSocketGateway({ namespace: 'events' })
-export class WsGateway implements OnGatewayInit, OnGatewayConnection {
+@WebSocketGateway({
+  cors: { origin: '*' },
+  transports: ['websocket'],
+})
+export class WsGateway implements OnGatewayConnection, OnModuleDestroy {
   @WebSocketServer()
   server: Server;
+  constructor(private tokenService: TokenService) {}
+
+  async handleConnection(client: Socket, ...args: any[]): Promise<void> {
+    try {
+      const token = await this.tokenService.verifyJwt(
+        client.handshake.auth?.token,
+      );
+      if (token.type !== JwtType.ACCESS) {
+        client.disconnect();
+      }
+    } catch (err) {
+      client.disconnect();
+    }
+  }
 
   @SubscribeMessage('message')
-  handleMessage(client: any, payload: any): string {
-    return 'Hello world!';
+  handleMessage(client: Socket, data: string): void {
+    client.broadcast.emit('message', data);
   }
 
-  handleConnection(client: any, ...args: any[]): any {
-    //
-  }
-
-  afterInit(server: any): any {
-    //
+  onModuleDestroy() {
+    if (this.server) {
+      this.server.close();
+    }
   }
 }

apps/server/src/ws/ws.module.ts

@@ -1,7 +1,9 @@
 import { Module } from '@nestjs/common';
 import { WsGateway } from './ws.gateway';
+import { AuthModule } from '../core/auth/auth.module';
 
 @Module({
+  imports: [AuthModule],
   providers: [WsGateway],
 })
 export class WsModule {}

package.json

@@ -17,6 +17,7 @@
     "@tiptap/extension-document": "^2.3.0",
     "@tiptap/extension-heading": "^2.3.0",
     "@tiptap/extension-highlight": "^2.3.0",
+    "@tiptap/extension-history": "^2.3.0",
     "@tiptap/extension-link": "^2.3.0",
     "@tiptap/extension-list-item": "^2.3.0",
     "@tiptap/extension-list-keymap": "^2.3.0",

pnpm-lock.yaml

@@ -41,6 +41,9 @@ importers:
       '@tiptap/extension-highlight':
         specifier: ^2.3.0
         version: 2.3.0(@tiptap/core@2.3.0)
+      '@tiptap/extension-history':
+        specifier: ^2.3.0
+        version: 2.3.0(@tiptap/core@2.3.0)(@tiptap/pm@2.3.0)
       '@tiptap/extension-link':
         specifier: ^2.3.0
         version: 2.3.0(@tiptap/core@2.3.0)(@tiptap/pm@2.3.0)