server.cjs

/**
 * Yjs WebSocket 服务器（基于官方最佳实践）
 * 参考：https://github.com/yjs/y-websocket
 */

const http = require("http");
const WebSocket = require("ws");
const Y = require("yjs");
const syncProtocol = require("y-protocols/sync");
const awarenessProtocol = require("y-protocols/awareness");
const encoding = require("lib0/encoding");
const decoding = require("lib0/decoding");

const PORT = process.env.PORT || 1234;

// 存储所有文档
const docs = new Map();

// 获取或创建 Yjs 文档
const getYDoc = (docname) => {
  let docData = docs.get(docname);
  if (!docData) {
    const doc = new Y.Doc();
    const awareness = new awarenessProtocol.Awareness(doc);
    awareness.setLocalState(null);

    // 确保初始化 objects Map（与客户端保持一致）
    doc.getMap("objects");

    docData = { doc, awareness, conns: new Map() };
    docs.set(docname, docData);

    console.log(`✓ 创建新文档: ${docname}`);
  }
  return docData;
};

// 发送消息（安全）
const send = (conn, message) => {
  if (conn.readyState === WebSocket.OPEN) {
    conn.send(message, (err) => {
      if (err) console.error("发送消息错误:", err.message);
    });
  }
};

/**
 * 官方推荐的 setupWSConnection 实现
 */
const setupWSConnection = (
  conn,
  req,
  docname = req.url.slice(1).split("?")[0]
) => {
  conn.binaryType = "arraybuffer";

  const { doc, awareness, conns } = getYDoc(docname);
  conns.set(conn, new Set());

  // 消息处理器
  const messageListener = (message) => {
    try {
      const encoder = encoding.createEncoder();
      const decoder = decoding.createDecoder(new Uint8Array(message));
      const messageType = decoding.readVarUint(decoder);

      switch (messageType) {
        case syncProtocol.messageYjsSyncStep1:
          encoding.writeVarUint(encoder, syncProtocol.messageYjsSyncStep2);
          syncProtocol.readSyncStep1(decoder, encoder, doc);
          send(conn, encoding.toUint8Array(encoder));
          break;

        case syncProtocol.messageYjsSyncStep2:
          syncProtocol.readSyncStep2(decoder, doc, conn);
          break;

        case syncProtocol.messageYjsUpdate:
          syncProtocol.readUpdate(decoder, doc, conn);
          break;

        case awarenessProtocol.messageAwareness:
          awarenessProtocol.applyAwarenessUpdate(
            awareness,
            decoding.readVarUint8Array(decoder),
            conn
          );
          break;

        default:
          console.warn("未知消息类型:", messageType);
      }
    } catch (err) {
      // 静默处理错误消息，避免日志污染
    }
  };

  // 文档更新监听器（广播给其他客户端）
  const updateListener = (update, origin) => {
    const encoder = encoding.createEncoder();
    encoding.writeVarUint(encoder, syncProtocol.messageYjsUpdate);
    encoding.writeVarUint8Array(encoder, update);
    const message = encoding.toUint8Array(encoder);

    conns.forEach((_, c) => {
      if (c !== origin) {
        send(c, message);
      }
    });
  };

  // Awareness 更新监听器（广播给所有客户端）
  const awarenessChangeListener = ({ added, updated, removed }, origin) => {
    const changedClients = added.concat(updated).concat(removed);
    const encoder = encoding.createEncoder();
    encoding.writeVarUint(encoder, awarenessProtocol.messageAwareness);
    encoding.writeVarUint8Array(
      encoder,
      awarenessProtocol.encodeAwarenessUpdate(awareness, changedClients)
    );
    const message = encoding.toUint8Array(encoder);

    // 广播给所有连接
    conns.forEach((_, c) => {
      send(c, message);
    });
  };

  // 连接关闭监听器
  const closeListener = () => {
    doc.off("update", updateListener);
    awareness.off("update", awarenessChangeListener);
    awareness.off("change", awarenessChangeListener);
    conns.delete(conn);

    // 清理 awareness 状态
    awarenessProtocol.removeAwarenessStates(
      awareness,
      Array.from(awareness.getStates().keys()).filter(
        (client) =>
          !Array.from(conns.keys()).some((c) => c.readyState === WebSocket.OPEN)
      ),
      null
    );

    console.log(`✗ 连接关闭: ${docname} (剩余: ${conns.size})`);

    // 如果没有连接了，5分钟后清理文档
    if (conns.size === 0) {
      setTimeout(() => {
        if (conns.size === 0) {
          docs.delete(docname);
          console.log(`♻️  清理文档: ${docname}`);
        }
      }, 5 * 60 * 1000);
    }
  };

  // 绑定所有监听器
  doc.on("update", updateListener);
  awareness.on("update", awarenessChangeListener);
  awareness.on("change", awarenessChangeListener);
  conn.on("message", messageListener);
  conn.on("close", closeListener);
  conn.on("error", closeListener);

  // 发送初始同步消息
  {
    const encoder = encoding.createEncoder();
    encoding.writeVarUint(encoder, syncProtocol.messageYjsSyncStep1);
    syncProtocol.writeSyncStep1(encoder, doc);
    send(conn, encoding.toUint8Array(encoder));
  }

  // 发送当前 awareness 状态
  if (awareness.getStates().size > 0) {
    const encoder = encoding.createEncoder();
    encoding.writeVarUint(encoder, awarenessProtocol.messageAwareness);
    encoding.writeVarUint8Array(
      encoder,
      awarenessProtocol.encodeAwarenessUpdate(
        awareness,
        Array.from(awareness.getStates().keys())
      )
    );
    send(conn, encoding.toUint8Array(encoder));
  }

  console.log(`✓ 新连接: ${docname}`);
};

// -------------------------
// HTTP + WebSocket 服务器
// -------------------------
const server = http.createServer((req, res) => {
  res.writeHead(200, { "Content-Type": "text/plain" });
  res.end("Yjs WebSocket Server\n");
});

const wss = new WebSocket.Server({ server });

wss.on("connection", setupWSConnection);

server.listen(PORT, () => {
  console.log(`✓ WebSocket 服务器运行在 ws://localhost:${PORT}`);
  console.log(`✓ HTTP 服务运行在 http://localhost:${PORT}`);
});

// 优雅关闭
process.on("SIGINT", () => {
  console.log("\n正在关闭服务器...");
  wss.close(() => {
    server.close(() => {
      console.log("服务器已关闭");
      process.exit(0);
    });
  });
});