Submatter/apps/collab-server/src/index.ts

import { createHash } from "node:crypto";
import { Forbidden } from "@hocuspocus/common";
import { Database } from "@hocuspocus/extension-database";
import { Redis as HocuspocusRedis } from "@hocuspocus/extension-redis";
import { Hocuspocus, type Extension, type onAuthenticatePayload } from "@hocuspocus/server";
import { db } from "@tasks/database/client";
import { objects } from "@tasks/database/schema";
import { eq } from "drizzle-orm";
import IoRedis from "ioredis";
import { z } from "zod";
import {
  base64ToYUpdate,
  extractObjectId,
  yDocToBase64,
} from "./utils.js";

const contentSchema = z
  .object({
    yjs: z.string().optional(),
  })
  .passthrough();

function parseYjsBase64FromContent(content: unknown): string | null {
  if (content == null) {
    return null;
  }
  if (typeof content === "string" && content.length > 0) {
    return content;
  }
  const parsed = contentSchema.safeParse(content);
  if (!parsed.success || !parsed.data.yjs) {
    return null;
  }
  return parsed.data.yjs;
}

async function tryCreateRedisExtension(): Promise<HocuspocusRedis | null> {
  const url = process.env.REDIS_URL?.trim();
  if (!url) {
    return null;
  }

  const client = new IoRedis(url, {
    maxRetriesPerRequest: null,
    enableReadyCheck: true,
    lazyConnect: false,
  });

  try {
    await client.ping();
  } catch (error) {
    console.warn(
      "[collab] Redis unavailable, continuing without Redis extension:",
      error,
    );
    try {
      client.disconnect();
    } catch {
      /* ignore */
    }
    return null;
  }

  return new HocuspocusRedis({
    redis: client,
    identifier: `collab-${process.pid}-${Date.now()}`,
  });
}

function buildDatabaseExtension() {
  return new Database({
    fetch: async ({ documentName, document }) => {
      const objectId = extractObjectId(documentName);
      if (!objectId) {
        console.error(
          `[collab] fetch: invalid document name (expected object:{id}): ${documentName}`,
        );
        return null;
      }

      try {
        const rows = await db
          .select({ content: objects.content })
          .from(objects)
          .where(eq(objects.id, objectId))
          .limit(1);

        const row = rows[0];
        if (!row) {
          return null;
        }

        const base64 = parseYjsBase64FromContent(row.content);
        if (!base64) {
          return null;
        }

        try {
          return base64ToYUpdate(base64);
        } catch (e) {
          console.error(
            `[collab] fetch: invalid Yjs base64 for object ${objectId}:`,
            e,
          );
          return null;
        }
      } catch (error) {
        console.error(
          `[collab] fetch: database error for object ${objectId}:`,
          error,
        );
        return null;
      }
    },

    store: async ({ documentName, document }) => {
      const objectId = extractObjectId(documentName);
      if (!objectId) {
        console.error(
          `[collab] store: invalid document name (expected object:{id}): ${documentName}`,
        );
        return;
      }

      const base64 = yDocToBase64(document);

      try {
        const existing = await db
          .select({ content: objects.content })
          .from(objects)
          .where(eq(objects.id, objectId))
          .limit(1);

        const prev = existing[0]?.content;
        const merged =
          prev !== null &&
          prev !== undefined &&
          typeof prev === "object" &&
          !Array.isArray(prev)
            ? { ...(prev as Record<string, unknown>), yjs: base64 }
            : { yjs: base64 };

        const updated = await db
          .update(objects)
          .set({
            content: merged,
            updatedAt: new Date(),
          })
          .where(eq(objects.id, objectId))
          .returning({ id: objects.id });

        if (updated.length === 0) {
          console.warn(
            `[collab] store: object ${objectId} not found, skipping persist`,
          );
        }
      } catch (error) {
        console.error(
          `[collab] store: database error for object ${objectId}:`,
          error,
        );
      }
    },
  });
}

const CURSOR_COLORS = [
  "#958DF1",
  "#F98181",
  "#FBBC88",
  "#FAF594",
  "#70CFF8",
  "#94FADB",
  "#B9F18D",
  "#E8A0BF",
];

function colorFromName(name: string): string {
  let hash = 0;
  for (let i = 0; i < name.length; i++) {
    hash = name.charCodeAt(i) + ((hash << 5) - hash);
  }
  return CURSOR_COLORS[Math.abs(hash) % CURSOR_COLORS.length];
}

async function main() {
  const port = Number(process.env.PORT) || 1234;

  const extensions: Extension[] = [buildDatabaseExtension()];
  const redis = await tryCreateRedisExtension();
  if (redis) {
    extensions.push(redis);
  }

  const hocuspocus = new Hocuspocus({
    quiet: true,
    stopOnSignals: false,
    address: "0.0.0.0",
    port,
    extensions,

    async onAuthenticate({ token }: onAuthenticatePayload) {
      if (!token?.trim()) {
        throw Forbidden;
      }

      // Try JWT decode for real user identity
      const authSecret = process.env.AUTH_SECRET;
      if (authSecret) {
        try {
          // Simple JWT decode (base64url decode the payload)
          const parts = token.split(".");
          if (parts.length === 3) {
            const payload = JSON.parse(
              Buffer.from(parts[1], "base64url").toString("utf-8"),
            );
            if (payload.sub || payload.name) {
              return {
                user: {
                  id: payload.sub ?? payload.email ?? token.slice(0, 16),
                  name: payload.name ?? payload.email ?? "Anonymous",
                  color: colorFromName(
                    payload.name ?? payload.email ?? "User",
                  ),
                },
              };
            }
          }
        } catch {
          // Fall through to hash-based identity
        }
      }

      const id = createHash("sha256").update(token).digest("hex");
      return {
        user: {
          id,
          name: `User ${id.slice(0, 8)}`,
          color: colorFromName(id),
        },
      };
    },

    async onConnect({ documentName, connection }) {
      const user = connection.readOnly ? "read-only" : "editor";
      console.log(`[collab] connect document=${documentName} (${user})`);
    },

    async onDisconnect({ documentName }) {
      console.log(`[collab] disconnect document=${documentName}`);
    },

    async afterUnloadDocument({ documentName }) {
      console.log(`[collab] afterUnloadDocument document=${documentName}`);
    },
  });

  const shutdown = async (signal: NodeJS.Signals) => {
    console.log(`[collab] received ${signal}, shutting down...`);
    try {
      await hocuspocus.destroy();
    } catch (error) {
      console.error("[collab] error during shutdown:", error);
    }
    process.exit(0);
  };

  process.once("SIGINT", () => void shutdown("SIGINT"));
  process.once("SIGTERM", () => void shutdown("SIGTERM"));

  await hocuspocus.listen();

  console.log(`Collaboration server running on port ${port}`);
}

main().catch((error) => {
  console.error("[collab] fatal error:", error);
  process.exit(1);
});