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

/**
 * Yjs WebSocket Server – Real-time collaboration with LevelDB persistence
 */
import * as Y from 'yjs';
import { LeveldbPersistence } from 'y-leveldb';
import type { FastifyInstance } from 'fastify';
import type { WebSocket } from '@fastify/websocket';

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

// Document store: docName → Y.Doc
const docs = new Map<string, Y.Doc>();
// 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> {
  try {
    persistence = new LeveldbPersistence(PERSISTENCE_DIR);
    console.log(`Yjs persistence initialized at ${PERSISTENCE_DIR}`);
  } catch (err) {
    console.warn('Yjs persistence failed to init (non-fatal):', err);
  }
}

export async function closeYjsPersistence(): Promise<void> {
  if (persistence) {
    await persistence.destroy();
    persistence = 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);

  // 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 broadcastUpdate(docName: string, update: Uint8Array, exclude?: WebSocket): void {
  const conns = connections.get(docName);
  if (!conns) return;

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

export function registerYjsWebSocket(fastify: FastifyInstance): 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;
    }

    const doc = await getOrCreateDoc(docName);

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

    // Send current document state to new client
    const stateUpdate = Y.encodeStateAsUpdate(doc);
    if (stateUpdate.length > 0 && socket.readyState === 1) {
      socket.send(stateUpdate);
    }

    // Listen for updates from this client
    socket.on('message', (data: Buffer) => {
      try {
        const update = new Uint8Array(data);
        Y.applyUpdate(doc, update);
        broadcastUpdate(docName, update, socket);
      } catch (err) {
        console.warn(`Invalid update from client for ${docName}:`, err);
      }
    });

    // Clean up on disconnect
    socket.on('close', () => {
      const conns = connections.get(docName);
      if (conns) {
        conns.delete(socket);
        if (conns.size === 0) {
          connections.delete(docName);
          // Optionally keep doc in memory for a while
          // For now, remove it to free memory
          docs.delete(docName);
        }
      }
    });

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