Next.js Tldraw Y.js Convex store

use-yjs-convex-store.ts

// https://github.com/tldraw/tldraw-yjs-example
import {
  InstancePresenceRecordType,
  TLAnyShapeUtilConstructor,
  TLInstancePresence,
  TLRecord,
  TLStoreWithStatus,
  computed,
  createPresenceStateDerivation,
  createTLStore,
  defaultShapeUtils,
  defaultUserPreferences,
  getUserPreferences,
  setUserPreferences,
  react,
  SerializedSchema,
  loadSnapshot,
  getSnapshot,
} from "tldraw";
import { useEffect, useMemo, useState } from "react";
import { YKeyValue } from "y-utility/y-keyvalue";
import { WebsocketProvider } from "y-websocket";
import * as Y from "yjs";
import { useUser } from "@clerk/nextjs";
import { api } from "@/convex/_generated/api";
import { useApiMutation } from "@/hooks/use-api-mutation";
import { useQuery } from "convex/react";
import { Id } from "@/convex/_generated/dataModel";

function throttle<T extends unknown[]>(
  func: (...args: T) => void,
  delay: number
): (...args: T) => void {
  let timerFlag: NodeJS.Timeout | null = null;

  return (...args: T) => {
    if (timerFlag === null) {
      func(...args);
      timerFlag = setTimeout(() => {
        timerFlag = null;
      }, delay);
    }
  };
}

export const useYjsStore = ({
  roomId,
  shapeUtils = [],
}: {
  roomId: string;
  shapeUtils?: TLAnyShapeUtilConstructor[];
}): TLStoreWithStatus => {
  // convex apis
  const { mutate } = useApiMutation(api.rooms.updateWhiteboardSnapshot);
  const data = useQuery(api.rooms.get, {
    id: roomId as Id<"rooms">,
  });

  // get user from clerk
  const { user } = useUser();

  // config ws
  const hostUrl =
    process.env.NODE_ENV === "development"
      ? `ws://localhost:3001`
      : "wss://demo.acadia.vn";

  // tldraw store
  const [tlStore] = useState(() => {
    const store = createTLStore({
      shapeUtils: [...defaultShapeUtils, ...shapeUtils],
    });
    return store;
  });

  // tl store status https://tldraw.dev/docs/persistence#The-store-prop
  const [storeWithStatus, setStoreWithStatus] = useState<TLStoreWithStatus>({
    status: "loading",
  });

  // yjs init
  const { yDoc, yStore, meta, room } = useMemo(() => {
    const yDoc = new Y.Doc({ gc: true });
    const yArr = yDoc.getArray<{ key: string; val: TLRecord }>(`tl_${roomId}`);
    const yStore = new YKeyValue(yArr);
    const meta = yDoc.getMap<SerializedSchema>("meta");

    return {
      yDoc,
      yStore,
      meta,
      room: new WebsocketProvider(hostUrl, roomId, yDoc, { connect: true }),
    };
  }, [roomId]);

  // load data from convex database
  useEffect(() => {
    if (storeWithStatus.status === "loading" && data) {
      // if there is whiteboardSnapshot load it to tldraw store
      if (data.whiteboardSnapshot) {
        const snapshot = JSON.parse(data.whiteboardSnapshot);
        loadSnapshot(tlStore, snapshot);
      }

      // tldraw store is ready
      setStoreWithStatus({
        store: tlStore,
        status: "synced-remote",
        connectionStatus: "online",
      });

      // when yjs updated, sync the changes to convex database every 5 seconds
      const onYDocUpdate = () => {
        const { document, session } = getSnapshot(tlStore);
        mutate({
          id: roomId,
          whiteboardSnapshot: JSON.stringify({ document, session }),
        });
      };
      yDoc.on("update", throttle(onYDocUpdate, 5000));
    }
  }, [storeWithStatus.status, data]);

  useEffect(() => {
    setStoreWithStatus({ status: "loading" });

    const unsubs: (() => void)[] = [];

    // Sync tldraw store changes to the yjs doc
    function handleSync() {
      unsubs.push(
        // https://tldraw.dev/docs/persistence#Listening-for-changes
        tlStore.listen(
          ({ changes: { added, updated, removed } }) => {
            yDoc.transact(() => {
              Object.values(added).forEach((record) =>
                yStore.set(record.id, record)
              );

              Object.values(updated).forEach(([_, record]) =>
                yStore.set(record.id, record)
              );

              Object.values(removed).forEach((record) =>
                yStore.delete(record.id)
              );
            });
          },
          { source: "user", scope: "document" } // only sync user's document changes
        )
      );

      // Sync the yjs doc changes to the tldraw store
      const handleChange = (
        changes: Map<
          string,
          | { action: "delete"; oldValue: TLRecord }
          | { action: "update"; oldValue: TLRecord; newValue: TLRecord }
          | { action: "add"; newValue: TLRecord }
        >,
        transaction: Y.Transaction
      ) => {
        if (transaction.local) return;

        const toRemove: TLRecord["id"][] = [];
        const toPut: TLRecord[] = [];

        changes.forEach((change, id) => {
          switch (change.action) {
            case "add":
            case "update": {
              const record = yStore.get(id)!;
              toPut.push(record);
              break;
            }
            case "delete": {
              toRemove.push(id as TLRecord["id"]);
              break;
            }
          }
        });

        // put / remove the records in the store
        tlStore.mergeRemoteChanges(() => {
          if (toRemove.length) tlStore.remove(toRemove);
          if (toPut.length) tlStore.put(toPut);
        });
      };

      yStore.on("change", handleChange);
      unsubs.push(() => yStore.off("change", handleChange));

      /* -------------------- Awareness ------------------- */

      const yClientId = room.awareness.clientID.toString();

      setUserPreferences({ id: yClientId, name: user?.fullName || "Guest" });

      const userPreferences = computed<{
        id: string;
        color: string;
        name: string;
      }>("userPreferences", () => {
        const user = getUserPreferences();
        return {
          id: user.id,
          color: user.color ?? defaultUserPreferences.color,
          name: user.name ?? "Guest",
        };
      });

      // Create the instance presence derivation
      const presenceId = InstancePresenceRecordType.createId(yClientId);
      const presenceDerivation = createPresenceStateDerivation(
        userPreferences,
        presenceId
      )(tlStore);

      // Set our initial presence from the derivation's current value
      room.awareness.setLocalStateField("presence", presenceDerivation.get());

      // When the derivation change, sync presence to to yjs awareness
      unsubs.push(
        react("when presence changes", () => {
          const presence = presenceDerivation.get();
          requestAnimationFrame(() => {
            room.awareness.setLocalStateField("presence", presence);
          });
        })
      );

      // Sync yjs awareness changes to the store
      const handleUpdate = (update: {
        added: number[];
        updated: number[];
        removed: number[];
      }) => {
        const states = room.awareness.getStates() as Map<
          number,
          { presence: TLInstancePresence }
        >;

        const toRemove: TLInstancePresence["id"][] = [];
        const toPut: TLInstancePresence[] = [];

        // Connect records to put / remove
        for (const clientId of update.added) {
          const state = states.get(clientId);
          if (state?.presence && state.presence.id !== presenceId) {
            toPut.push(state.presence);
          }
        }

        for (const clientId of update.updated) {
          const state = states.get(clientId);
          if (state?.presence && state.presence.id !== presenceId) {
            toPut.push(state.presence);
          }
        }

        for (const clientId of update.removed) {
          toRemove.push(
            InstancePresenceRecordType.createId(clientId.toString())
          );
        }

        // put / remove the records in the store
        tlStore.mergeRemoteChanges(() => {
          if (toRemove.length) tlStore.remove(toRemove);
          if (toPut.length) tlStore.put(toPut);
        });
      };

      const handleMetaUpdate = () => {
        const theirSchema = meta.get("schema");
        if (!theirSchema) {
          throw new Error("No schema found in the yjs doc");
        }
        // If the shared schema is newer than our schema, the user must refresh
        const newMigrations = tlStore.schema.getMigrationsSince(theirSchema);

        if (!newMigrations.ok || newMigrations.value.length > 0) {
          window.alert("The schema has been updated. Please refresh the page.");
          yDoc.destroy();
        }
      };
      meta.observe(handleMetaUpdate);
      unsubs.push(() => meta.unobserve(handleMetaUpdate));

      room.awareness.on("update", handleUpdate);
      unsubs.push(() => room.awareness.off("update", handleUpdate));
    }

    let hasConnectedBefore = false;

    function handleStatusChange({
      status,
    }: {
      status: "disconnected" | "connected" | "connecting";
    }) {
      // If we're disconnected, set the store status to 'synced-remote' and the connection status to 'offline'
      if (status === "disconnected") {
        setStoreWithStatus({
          store: tlStore,
          status: "synced-remote",
          connectionStatus: "offline",
        });
        return;
      }

      room.off("sync", handleSync);

      if (status === "connected") {
        if (hasConnectedBefore) return;
        hasConnectedBefore = true;
        room.on("sync", handleSync);
        unsubs.push(() => room.off("sync", handleSync));
      }
    }

    room.on("status", handleStatusChange);

    unsubs.push(() => room.off("status", handleStatusChange));

    return () => {
      unsubs.forEach((fn) => fn());
      unsubs.length = 0;
    };
  }, [room, yDoc, tlStore, yStore, meta]);

  return storeWithStatus;
};