Merge pull request #48 from trails-cool/fix-yjs-sync-protocol

apps/planner/app/lib/yjs-server.ts

@@ -1,10 +1,18 @@
 import { WebSocketServer, type WebSocket } from "ws";
 import * as Y from "yjs";
+import * as syncProtocol from "y-protocols/sync";
+import * as awarenessProtocol from "y-protocols/awareness";
+import * as encoding from "lib0/encoding";
+import * as decoding from "lib0/decoding";
 import type { IncomingMessage, Server } from "node:http";
 import { saveSessionState, loadSessionState, touchSession } from "./sessions.ts";
 
+const messageSync = 0;
+const messageAwareness = 1;
+
 const docs = new Map<string, Y.Doc>();
-const sessionClients = new Map<string, Set<WebSocket>>();
+const awarenessMap = new Map<string, awarenessProtocol.Awareness>();
+const conns = new Map<WebSocket, { sessionId: string; clientIds: Set<number> }>();
 
 export function getOrCreateDoc(sessionId: string): Y.Doc {
   let doc = docs.get(sessionId);
@@ -28,12 +36,39 @@ export async function getOrLoadDoc(sessionId: string): Promise<Y.Doc> {
   return doc;
 }
 
+function getAwareness(sessionId: string, doc: Y.Doc): awarenessProtocol.Awareness {
+  let awareness = awarenessMap.get(sessionId);
+  if (!awareness) {
+    awareness = new awarenessProtocol.Awareness(doc);
+    awarenessMap.set(sessionId, awareness);
+
+    awareness.on("update", ({ added, updated, removed }: { added: number[]; updated: number[]; removed: number[] }) => {
+      const changedClients = added.concat(updated, removed);
+      const encoder = encoding.createEncoder();
+      encoding.writeVarUint(encoder, messageAwareness);
+      encoding.writeVarUint8Array(encoder, awarenessProtocol.encodeAwarenessUpdate(awareness!, changedClients));
+      const message = encoding.toUint8Array(encoder);
+
+      for (const [ws, meta] of conns) {
+        if (meta.sessionId === sessionId && ws.readyState === ws.OPEN) {
+          ws.send(message);
+        }
+      }
+    });
+  }
+  return awareness;
+}
+
 export function deleteDoc(sessionId: string): boolean {
   const doc = docs.get(sessionId);
   if (doc) {
     doc.destroy();
     docs.delete(sessionId);
-    sessionClients.delete(sessionId);
+    const awareness = awarenessMap.get(sessionId);
+    if (awareness) {
+      awareness.destroy();
+      awarenessMap.delete(sessionId);
+    }
     return true;
   }
   return false;
@@ -44,7 +79,11 @@ export function getDocCount(): number {
 }
 
 export function getClientCount(sessionId: string): number {
-  return sessionClients.get(sessionId)?.size ?? 0;
+  let count = 0;
+  for (const meta of conns.values()) {
+    if (meta.sessionId === sessionId) count++;
+  }
+  return count;
 }
 
 // Debounced persistence — save at most every 5 seconds per session
@@ -65,6 +104,43 @@ function debouncedSave(sessionId: string): void {
   );
 }
 
+function handleMessage(ws: WebSocket, message: Uint8Array, sessionId: string) {
+  const doc = docs.get(sessionId);
+  if (!doc) return;
+
+  const awareness = awarenessMap.get(sessionId);
+  const decoder = decoding.createDecoder(message);
+  const messageType = decoding.readVarUint(decoder);
+
+  switch (messageType) {
+    case messageSync: {
+      const encoder = encoding.createEncoder();
+      encoding.writeVarUint(encoder, messageSync);
+      syncProtocol.readSyncMessage(decoder, encoder, doc, ws);
+      const resp = encoding.toUint8Array(encoder);
+      if (encoding.length(encoder) > 1) {
+        ws.send(resp);
+      }
+      debouncedSave(sessionId);
+      break;
+    }
+    case messageAwareness: {
+      if (awareness) {
+        const update = decoding.readVarUint8Array(decoder);
+        awarenessProtocol.applyAwarenessUpdate(awareness, update, ws);
+
+        const meta = conns.get(ws);
+        if (meta) {
+          awareness.getStates().forEach((_state, clientId) => {
+            meta.clientIds.add(clientId);
+          });
+        }
+      }
+      break;
+    }
+  }
+}
+
 export function setupYjsWebSocket(server: Server): WebSocketServer {
   const wss = new WebSocketServer({ noServer: true });
 
@@ -84,50 +160,70 @@ export function setupYjsWebSocket(server: Server): WebSocketServer {
 
   wss.on("connection", async (ws: WebSocket, _request: IncomingMessage, sessionId: string) => {
     const doc = await getOrLoadDoc(sessionId);
+    const awareness = getAwareness(sessionId, doc);
 
-    // Track clients per session
-    if (!sessionClients.has(sessionId)) {
-      sessionClients.set(sessionId, new Set());
+    conns.set(ws, { sessionId, clientIds: new Set() });
+
+    // Broadcast doc updates to all connections in this session
+    const onUpdate = (update: Uint8Array, origin: unknown) => {
+      if (origin === ws) return;
+      const encoder = encoding.createEncoder();
+      encoding.writeVarUint(encoder, messageSync);
+      syncProtocol.writeUpdate(encoder, update);
+      const message = encoding.toUint8Array(encoder);
+
+      for (const [client, meta] of conns) {
+        if (meta.sessionId === sessionId && client !== ws && client.readyState === ws.OPEN) {
+          client.send(message);
+        }
+      }
+    };
+    doc.on("update", onUpdate);
+
+    // Send sync step 1
+    const encoder = encoding.createEncoder();
+    encoding.writeVarUint(encoder, messageSync);
+    syncProtocol.writeSyncStep1(encoder, doc);
+    ws.send(encoding.toUint8Array(encoder));
+
+    // Send current awareness states
+    const awarenessStates = awareness.getStates();
+    if (awarenessStates.size > 0) {
+      const awarenessEncoder = encoding.createEncoder();
+      encoding.writeVarUint(awarenessEncoder, messageAwareness);
+      encoding.writeVarUint8Array(
+        awarenessEncoder,
+        awarenessProtocol.encodeAwarenessUpdate(awareness, Array.from(awarenessStates.keys())),
+      );
+      ws.send(encoding.toUint8Array(awarenessEncoder));
     }
-    sessionClients.get(sessionId)!.add(ws);
-
-    // Send current state to new client
-    const state = Y.encodeStateAsUpdate(doc);
-    ws.send(state);
 
     await touchSession(sessionId);
 
-    // Listen for updates from this client
-    ws.on("message", (data: Buffer) => {
+    ws.on("message", (data: ArrayBuffer) => {
       try {
-        const update = new Uint8Array(data);
-        Y.applyUpdate(doc, update);
-
-        // Broadcast to all other clients in this session
-        const clients = sessionClients.get(sessionId);
-        if (clients) {
-          for (const client of clients) {
-            if (client !== ws && client.readyState === ws.OPEN) {
-              client.send(data);
-            }
-          }
-        }
-
-        // Persist state (debounced)
-        debouncedSave(sessionId);
+        handleMessage(ws, new Uint8Array(data), sessionId);
       } catch (_e) {
-        // Ignore malformed updates
+        // Ignore malformed messages
       }
     });
 
     ws.on("close", () => {
-      const clients = sessionClients.get(sessionId);
-      if (clients) {
-        clients.delete(ws);
-        if (clients.size === 0) {
-          // Last client left — save immediately
-          saveSessionState(sessionId).catch(() => {});
-        }
+      const meta = conns.get(ws);
+      if (meta && meta.clientIds.size > 0) {
+        awarenessProtocol.removeAwarenessStates(
+          awareness,
+          Array.from(meta.clientIds),
+          null,
+        );
+      }
+      conns.delete(ws);
+      doc.off("update", onUpdate);
+
+      // Save when last client leaves
+      const hasClients = Array.from(conns.values()).some((m) => m.sessionId === sessionId);
+      if (!hasClients) {
+        saveSessionState(sessionId).catch(() => {});
       }
     });
   });