From cc041d3235eb12cf2f4e7bb6d0450301c9d8a08b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ullrich=20Sch=C3=A4fer?= Date: Wed, 25 Mar 2026 02:24:13 +0100 Subject: [PATCH] Fix Yjs WebSocket server to use sync protocol The production server was sending raw Yjs updates but the client (y-websocket's WebsocketProvider) expects the y-protocols sync protocol with message type headers. This caused "Unexpected end of array" errors on every connection. Rewrite to match the Vite dev plugin: use y-protocols/sync for document sync, y-protocols/awareness for presence, and lib0 for encoding. Verified locally: server starts, WebSocket connects, sync messages exchange correctly. Co-Authored-By: Claude Opus 4.6 (1M context) --- apps/planner/app/lib/yjs-server.ts | 168 ++++++++++++++++++++++------- 1 file changed, 132 insertions(+), 36 deletions(-) diff --git a/apps/planner/app/lib/yjs-server.ts b/apps/planner/app/lib/yjs-server.ts index 50bc9d2..ccb22ba 100644 --- a/apps/planner/app/lib/yjs-server.ts +++ b/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(); -const sessionClients = new Map>(); +const awarenessMap = new Map(); +const conns = new Map }>(); export function getOrCreateDoc(sessionId: string): Y.Doc { let doc = docs.get(sessionId); @@ -28,12 +36,39 @@ export async function getOrLoadDoc(sessionId: string): Promise { 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(() => {}); } }); });