web-cad/backend/src/websocket/yjsServer.ts

/**
 * Yjs WebSocket Server – Real-time collaboration with LevelDB persistence.
 *
 * Message protocol (byte-prefixed):
 *   0 = Y.Doc update (Y.encodeStateAsUpdate or Y.encodeStateAsUpdateYUpdate)
 *   1 = Awareness update (encodeAwarenessUpdate from y-protocols/awareness)
 */
import * as Y from 'yjs';
import { LeveldbPersistence } from 'y-leveldb';
import { Awareness, encodeAwarenessUpdate, applyAwarenessUpdate, removeAwarenessStates } from 'y-protocols/awareness';
import type { FastifyInstance } from 'fastify';
import type { WebSocket } from '@fastify/websocket';
import type { AuthService } from '../auth/AuthService.js';

const PERSISTENCE_DIR = process.env.YJS_PERSISTENCE_DIR || '/tmp/yjs-documents';

// Message type prefixes (must match frontend WebSocketProvider.ts)
const MSG_DOC_UPDATE = 0;
const MSG_AWARENESS = 1;

// Document store: docName → Y.Doc
const docs = new Map<string, Y.Doc>();
// Awareness store: docName → Awareness
const awarenessMap = new Map<string, Awareness>();
// Connection tracking: docName → Set<WebSocket>
const connections = new Map<string, Set<WebSocket>>();
let persistence: LeveldbPersistence | null = null;
let persistenceReady: Promise<void> | null = null;

export async function initYjsPersistence(): Promise<void> {
  const tryInit = async (allowReset: boolean): Promise<void> => {
    try {
      persistence = new LeveldbPersistence(PERSISTENCE_DIR);
      persistenceReady = (async () => {
        // Wait for LevelDB to be ready
        await (persistence as any).getAllDocNames();
      })();
      await persistenceReady;
      console.log(`Yjs persistence initialized at ${PERSISTENCE_DIR}`);
    } catch (err) {
      // In dev mode the persistence dir lives under /tmp. If the previous process
      // did not shut down cleanly, LevelDB can be left in a locked/corrupted state.
      // Resetting the directory is acceptable because collaboration history is ephemeral.
      if (allowReset && PERSISTENCE_DIR.startsWith('/tmp/')) {
        console.warn(`Yjs persistence init failed, resetting ${PERSISTENCE_DIR} and retrying...`);
        try {
          if (persistence) {
            await persistence.destroy().catch(() => {});
            persistence = null;
            persistenceReady = null;
          }
          const fs = await import('fs');
          fs.rmSync(PERSISTENCE_DIR, { recursive: true, force: true });
          fs.mkdirSync(PERSISTENCE_DIR, { recursive: true });
        } catch (cleanupErr) {
          console.warn('Yjs persistence cleanup failed:', cleanupErr);
        }
        return tryInit(false);
      }
      console.warn('Yjs persistence failed to init (running without disk persistence):', err);
      persistence = null;
      persistenceReady = null;
    }
  };

  await tryInit(true);
}

export async function closeYjsPersistence(): Promise<void> {
  if (persistence) {
    try {
      await persistence.destroy();
    } catch (err) {
      console.warn('Yjs persistence close failed (non-fatal):', err);
    }
    persistence = null;
  }
  persistenceReady = null;
}

async function getOrCreateDoc(docName: string): Promise<Y.Doc> {
  if (docs.has(docName)) {
    return docs.get(docName)!;
  }

  const doc = new Y.Doc();
  docs.set(docName, doc);

  // Create awareness instance for this document
  const awareness = new Awareness(doc);
  awarenessMap.set(docName, awareness);

  // Load persisted state if available
  if (persistence && persistenceReady) {
    try {
      await persistenceReady;
      const persistedYdoc = await persistence.getYDoc(docName);
      const update = Y.encodeStateAsUpdate(persistedYdoc);
      if (update.length > 0) {
        Y.applyUpdate(doc, update);
      }
    } catch (err) {
      console.warn(`Failed to load doc ${docName}:`, err);
    }
  }

  // Persist updates to LevelDB
  doc.on('update', (update: Uint8Array) => {
    if (persistence) {
      persistence.storeUpdate(docName, update).catch((err: any) => {
        console.warn(`Failed to persist update for ${docName}:`, err);
      });
    }
  });

  return doc;
}

function getOrCreateAwareness(docName: string): Awareness | null {
  return awarenessMap.get(docName) ?? null;
}

function broadcastUpdate(docName: string, update: Uint8Array, msgType: number, exclude?: WebSocket): void {
  const conns = connections.get(docName);
  if (!conns) return;

  // Prepend message type byte
  const msg = new Uint8Array(update.length + 1);
  msg[0] = msgType;
  msg.set(update, 1);

  for (const ws of conns) {
    if (ws !== exclude && ws.readyState === 1 /* OPEN */) {
      try {
        ws.send(msg);
      } catch {
        // Connection error — will be cleaned up on close
      }
    }
  }
}

export function registerYjsWebSocket(fastify: FastifyInstance, authService: AuthService): void {
  fastify.get('/ws/collab/:docName', { websocket: true }, async (socket: WebSocket, request: any) => {
    const docName = request.params.docName as string;
    if (!docName) {
      socket.close(4000, 'Missing docName');
      return;
    }

    // Issue #2: Authenticate WebSocket connection via query param token
    const token = request.query.token as string;
    if (!token) {
      socket.close(4001, 'Authentication required');
      return;
    }
    const user = authService.getUserFromSession(token);
    if (!user) {
      socket.close(4001, 'Invalid or expired session');
      return;
    }

    const doc = await getOrCreateDoc(docName);
    const awareness = getOrCreateAwareness(docName);

    // Track connection
    if (!connections.has(docName)) {
      connections.set(docName, new Set());
    }
    connections.get(docName)!.add(socket);

    // Send current document state to new client (with MSG_DOC_UPDATE prefix)
    const stateUpdate = Y.encodeStateAsUpdate(doc);
    if (stateUpdate.length > 0 && socket.readyState === 1) {
      const msg = new Uint8Array(stateUpdate.length + 1);
      msg[0] = MSG_DOC_UPDATE;
      msg.set(stateUpdate, 1);
      socket.send(msg);
    }

    // Send current awareness states to new client
    if (awareness && awareness.getStates().size > 0 && socket.readyState === 1) {
      const awarenessUpdate = encodeAwarenessUpdate(awareness, Array.from(awareness.getStates().keys()));
      if (awarenessUpdate.length > 0) {
        const msg = new Uint8Array(awarenessUpdate.length + 1);
        msg[0] = MSG_AWARENESS;
        msg.set(awarenessUpdate, 1);
        socket.send(msg);
      }
    }

    // Listen for messages from this client
    socket.on('message', (data: Buffer) => {
      try {
        const raw = new Uint8Array(data);
        if (raw.length === 0) return;

        const msgType = raw[0];
        const payload = raw.slice(1);

        if (msgType === MSG_AWARENESS) {
          // Awareness update — apply to server awareness and broadcast
          if (awareness) {
            applyAwarenessUpdate(awareness, payload, socket);
            broadcastUpdate(docName, payload, MSG_AWARENESS, socket);
          }
        } else if (msgType === MSG_DOC_UPDATE) {
          // Y.Doc update — apply to server doc and broadcast
          Y.applyUpdate(doc, payload);
          broadcastUpdate(docName, payload, MSG_DOC_UPDATE, socket);
        } else {
          // Legacy: treat entire message as a Y.Doc update (no prefix)
          Y.applyUpdate(doc, raw);
          broadcastUpdate(docName, raw, MSG_DOC_UPDATE, socket);
        }
      } catch (err) {
        console.warn(`Invalid update from client for ${docName}:`, err);
      }
    });

    // Clean up on disconnect (Issue #5: keep doc in memory for reconnection)
    socket.on('close', () => {
      const conns = connections.get(docName);
      if (conns) {
        conns.delete(socket);

        // Remove this client's awareness state
        if (awareness) {
          // Find and remove the client's awareness state based on socket origin
          const statesToRemove: number[] = [];
          for (const [clientID, meta] of awareness.meta) {
            if (meta === socket) {
              statesToRemove.push(clientID);
            }
          }
          if (statesToRemove.length > 0) {
            removeAwarenessStates(awareness, statesToRemove, socket);
            // Broadcast removal to remaining clients
            const removalUpdate = encodeAwarenessUpdate(awareness, statesToRemove);
            if (removalUpdate.length > 0) {
              broadcastUpdate(docName, removalUpdate, MSG_AWARENESS, socket);
            }
          }
        }

        if (conns.size === 0) {
          connections.delete(docName);
          // Keep doc in memory for reconnection (persistence handles durability)
        }
      }
    });

    socket.on('error', () => {
      const conns = connections.get(docName);
      if (conns) {
        conns.delete(socket);
      }
    });
  });
}