colabbd/src/session.ts

import type { ServerWebSocket } from "bun";
import * as Y from "yjs";
import { encode } from "./protocol";

export interface WsData {
  room: string | null;
  client: Client | null;
}

export interface Client {
  ws: ServerWebSocket<WsData>;
}

export class Session {
  doc: Y.Doc;
  clients: Set<Client> = new Set();

  constructor(public name: string) {
    this.doc = new Y.Doc();
  }

  join(client: Client) {
    this.clients.add(client);
    // send full state to new client
    const state = Y.encodeStateAsUpdate(this.doc);
    try {
      client.ws.send(encode({ type: "sync", data: Array.from(state) }));
    } catch (err) {
      console.debug("failed to send sync to client, removing:", err);
      this.clients.delete(client);
      return;
    }
    this.broadcastPeerCount();
  }

  leave(client: Client) {
    this.clients.delete(client);
    this.broadcastPeerCount();
    if (this.isEmpty()) {
      sessions.delete(this.name);
      console.debug(`session removed: ${this.name}`);
    }
  }

  isEmpty(): boolean {
    return this.clients.size === 0;
  }

  applyUpdate(update: Uint8Array, from: Client) {
    try {
      Y.applyUpdate(this.doc, update);
      // broadcast to others
      for (const client of this.clients) {
        if (client !== from) {
          try {
            client.ws.send(
              encode({ type: "update", data: Array.from(update) }),
            );
          } catch (err) {
            console.debug("failed to send update to client, removing:", err);
            this.clients.delete(client);
          }
        }
      }
    } catch (err) {
      console.error(`failed to apply update in session ${this.name}:`, err);
      // optionally notify the client that sent the bad update
      try {
        from.ws.send(encode({ type: "error", message: "invalid update" }));
      } catch {
        // ignore send errors
      }
    }
  }

  broadcastPeerCount() {
    const msg = encode({ type: "peers", count: this.clients.size });
    for (const client of this.clients) {
      try {
        client.ws.send(msg);
      } catch (err) {
        console.debug("failed to send peer count to client, removing:", err);
        this.clients.delete(client);
      }
    }
  }
}

// room name -> session
const sessions = new Map<string, Session>();

export function getOrCreateSession(name: string): Session {
  let session = sessions.get(name);
  if (!session) {
    session = new Session(name);
    sessions.set(name, session);
    console.debug(`session created: ${name}`);
  }
  return session;
}

export function getSession(name: string): Session | undefined {
  return sessions.get(name);
}

export function removeSession(name: string): void {
  const session = sessions.get(name);
  if (session?.isEmpty()) {
    sessions.delete(name);
    console.debug(`session removed: ${name}`);
  }
}