DonutBot/server.js

"use strict";

/**
 * Standalone control server.
 *
 * - Serves web/index.html at GET /
 * - Browser WebSocket clients connect at  /ws
 * - afk_parent connects as a single privileged client at  /parent
 *   (optionally authenticated via ?secret=... query param when
 *    PARENT_SECRET env var is set)
 *
 * The server relays messages:
 *   browser → parent  (bot.start / bot.stop / bot.chat / window.* / state.get)
 *   parent  → browsers (bot.update / bot.log / bot.event / state.full / ack)
 *
 * Each browser request gets a _relayId tag so the corresponding ack/response
 * from the parent can be routed back to the exact browser that made the call.
 */

const http = require("http");
const fs = require("fs");
const path = require("path");
const { WebSocketServer, WebSocket } = require("ws");

const WEB_PORT = Number(process.env.AFK_PARENT_PORT || 3008);
const PARENT_SECRET = process.env.PARENT_SECRET || "";
const WEB_UI_PATH = path.join(__dirname, "web", "index.html");

// ─── state ───────────────────────────────────────────────────────────────────

const parentsById = new Map(); // parentId -> ws
const wsToParentId = new Map(); // ws -> parentId
const browsers = new Map(); // sessionId → ws
let sessionCounter = 0;
let relayCounter = 0;
const pendingRelays = new Map(); // relayId → { sessionId, originalRequestId, parentId }
const cachedBots = new Map(); // `${parentId}:${username}` -> bot

// ─── helpers ─────────────────────────────────────────────────────────────────

function nowIso() {
  return new Date().toISOString();
}

function sendTo(ws, payload) {
  if (!ws || ws.readyState !== WebSocket.OPEN) return;
  try {
    ws.send(JSON.stringify(payload));
  } catch (_) {
    /* ignore */
  }
}

function broadcastToBrowsers(payload) {
  for (const ws of browsers.values()) sendTo(ws, payload);
}

function botKey(parentId, username) {
  return `${parentId}:${username}`;
}

function listCachedBots() {
  return [...cachedBots.values()];
}

// ─── forward server console logs to browsers ─────────────────────────────────
{
  const _log = console.log.bind(console);
  const _err = console.error.bind(console);
  const strip = (s) => s.replace(/^\[server\] /, "");
  console.log = (...args) => {
    const line = args.map(String).join(" ");
    _log(line);
    broadcastToBrowsers({
      type: "server.log",
      serverTime: nowIso(),
      stream: "stdout",
      line: strip(line),
    });
  };
  console.error = (...args) => {
    const line = args.map(String).join(" ");
    _err(line);
    broadcastToBrowsers({
      type: "server.log",
      serverTime: nowIso(),
      stream: "stderr",
      line: strip(line),
    });
  };
}

function nextRelayId(sessionId, originalRequestId, parentId) {
  relayCounter++;
  const id = `relay-${relayCounter}`;
  pendingRelays.set(id, { sessionId, originalRequestId, parentId });
  return id;
}

// ─── HTTP server ──────────────────────────────────────────────────────────────

const httpServer = http.createServer((req, res) => {
  if (req.method !== "GET") {
    res.writeHead(405, { "Content-Type": "text/plain" });
    res.end("Method not allowed");
    return;
  }

  if (req.url === "/" || req.url === "/index.html") {
    try {
      const html = fs.readFileSync(WEB_UI_PATH, "utf8");
      res.writeHead(200, { "Content-Type": "text/html; charset=utf-8" });
      res.end(html);
    } catch {
      res.writeHead(500, { "Content-Type": "text/plain" });
      res.end("UI file missing");
    }
    return;
  }

  if (req.url === "/health") {
    res.writeHead(200, { "Content-Type": "application/json; charset=utf-8" });
    res.end(
      JSON.stringify({
        ok: true,
        parentConnected: parentsById.size > 0,
        parentCount: parentsById.size,
        bots: listCachedBots(),
      }),
    );
    return;
  }

  res.writeHead(404, { "Content-Type": "text/plain" });
  res.end("Not found");
});

// ─── WebSocket upgrade routing ────────────────────────────────────────────────

const browserWss = new WebSocketServer({ noServer: true });
const parentWssI = new WebSocketServer({ noServer: true });

httpServer.on("upgrade", (req, socket, head) => {
  let pathname;
  try {
    pathname = new URL(req.url, "http://x").pathname;
  } catch {
    socket.destroy();
    return;
  }

  if (pathname === "/parent") {
    if (PARENT_SECRET) {
      let secret;
      try {
        secret = new URL(req.url, "http://x").searchParams.get("secret");
      } catch {
        /* ignore */
      }
      if (secret !== PARENT_SECRET) {
        socket.write("HTTP/1.1 401 Unauthorized\r\nConnection: close\r\n\r\n");
        socket.destroy();
        return;
      }
    }
    parentWssI.handleUpgrade(req, socket, head, (ws) =>
      parentWssI.emit("connection", ws, req),
    );
  } else if (pathname === "/ws") {
    browserWss.handleUpgrade(req, socket, head, (ws) =>
      browserWss.emit("connection", ws),
    );
  } else {
    socket.destroy();
  }
});

// ─── parent connection ────────────────────────────────────────────────────────

parentWssI.on("connection", (ws, req) => {
  const parentId =
    new URL(req.url, "http://x").searchParams.get("parentId") ||
    `parent-${Date.now()}-${Math.floor(Math.random() * 10_000)}`;

  const existing = parentsById.get(parentId);
  if (existing && existing !== ws) {
    console.log(
      `[server] replacing existing parent connection for ${parentId}`,
    );
    existing.terminate();
  }

  parentsById.set(parentId, ws);
  wsToParentId.set(ws, parentId);

  console.log(`[server] parent connected (${parentId})`);
  broadcastToBrowsers({
    type: "parent.connected",
    serverTime: nowIso(),
    parentId,
    parentCount: parentsById.size,
  });

  ws.on("message", (raw) => {
    let msg;
    try {
      msg = JSON.parse(String(raw));
    } catch {
      return;
    }
    handleParentMessage(parentId, msg);
  });

  ws.on("close", () => {
    const closedParentId = wsToParentId.get(ws);
    wsToParentId.delete(ws);
    if (!closedParentId) return;

    const activeWs = parentsById.get(closedParentId);
    if (activeWs === ws) parentsById.delete(closedParentId);

    for (const key of [...cachedBots.keys()]) {
      if (key.startsWith(`${closedParentId}:`)) cachedBots.delete(key);
    }

    console.log(`[server] parent disconnected (${closedParentId})`);
    broadcastToBrowsers({
      type: "parent.disconnected",
      serverTime: nowIso(),
      parentId: closedParentId,
      parentCount: parentsById.size,
    });
    broadcastToBrowsers({
      type: "state.full",
      serverTime: nowIso(),
      bots: listCachedBots(),
      parentConnected: parentsById.size > 0,
      parentCount: parentsById.size,
    });
  });

  ws.on("error", (err) => {
    console.error("[server] parent ws error:", err.message);
  });
});

function handleParentMessage(parentId, msg) {
  if (!msg || typeof msg !== "object") return;

  // ── broadcasts ──────────────────────────────────────────────────────────────

  if (msg.type === "bot.update") {
    if (msg.bot) {
      const bot = { ...msg.bot, parentId };
      cachedBots.set(botKey(parentId, bot.username), bot);
    }
    broadcastToBrowsers({
      ...msg,
      bot: msg.bot ? { ...msg.bot, parentId } : msg.bot,
      parentId,
      serverTime: msg.serverTime || nowIso(),
    });
    return;
  }

  if (
    msg.type === "bot.log" ||
    msg.type === "bot.event" ||
    msg.type === "parent.log"
  ) {
    broadcastToBrowsers({
      ...msg,
      parentId,
      serverTime: msg.serverTime || nowIso(),
    });
    return;
  }

  // ── state snapshot (response to state.get, or initial push on connect) ──────

  if (msg.type === "state.full") {
    if (Array.isArray(msg.bots)) {
      for (const bot of msg.bots) {
        if (!bot || !bot.username) continue;
        cachedBots.set(botKey(parentId, bot.username), { ...bot, parentId });
      }
    }

    const relay = msg._relayId ? pendingRelays.get(msg._relayId) : null;
    if (relay) {
      if (relay.parentId && relay.parentId !== parentId) return;
      pendingRelays.delete(msg._relayId);
      const ws = browsers.get(relay.sessionId);
      if (ws)
        sendTo(ws, {
          ...msg,
          parentId,
          bots: (msg.bots || []).map((bot) => ({ ...bot, parentId })),
          requestId: relay.originalRequestId,
          _relayId: undefined,
        });
    } else {
      broadcastToBrowsers({
        ...msg,
        parentId,
        bots: listCachedBots(),
        serverTime: msg.serverTime || nowIso(),
        _relayId: undefined,
        parentConnected: parentsById.size > 0,
        parentCount: parentsById.size,
      });
    }
    return;
  }

  // ── ack for a specific browser action ───────────────────────────────────────

  if (msg.type === "ack" && msg._relayId) {
    const relay = pendingRelays.get(msg._relayId);
    if (!relay) return;
    if (relay.parentId && relay.parentId !== parentId) return;
    pendingRelays.delete(msg._relayId);
    const ws = browsers.get(relay.sessionId);
    if (ws)
      sendTo(ws, {
        ...msg,
        parentId,
        requestId: relay.originalRequestId,
        _relayId: undefined,
      });
  }
}

// ─── browser connections ──────────────────────────────────────────────────────

browserWss.on("connection", (ws) => {
  sessionCounter++;
  const sessionId = `s${sessionCounter}`;
  browsers.set(sessionId, ws);

  // Send initial state immediately from cache; browser shows "parent offline"
  // badge when parentConnected is false.
  sendTo(ws, {
    type: "hello",
    serverTime: nowIso(),
    bots: listCachedBots(),
    parentConnected: parentsById.size > 0,
    parentCount: parentsById.size,
  });

  ws.on("message", (raw) => {
    let msg;
    try {
      msg = JSON.parse(String(raw));
    } catch {
      sendTo(ws, {
        type: "ack",
        requestId: null,
        ok: false,
        error: "Invalid JSON",
      });
      return;
    }
    handleBrowserMessage(sessionId, ws, msg);
  });

  ws.on("close", () => {
    browsers.delete(sessionId);
    // Remove any pending relays so we don't leak memory
    for (const [id, relay] of pendingRelays) {
      if (relay.sessionId === sessionId) pendingRelays.delete(id);
    }
  });
});

function handleBrowserMessage(sessionId, ws, msg) {
  if (!msg || typeof msg !== "object") return;

  if (parentsById.size === 0) {
    sendTo(ws, {
      type: "ack",
      requestId: msg.requestId || null,
      ok: false,
      error: "Parent not connected",
    });
    return;
  }

  let targetParentId = msg.parentId || null;

  if (!targetParentId && msg.username) {
    const parentIds = new Set();
    for (const bot of cachedBots.values()) {
      if (bot.username === msg.username && bot.parentId)
        parentIds.add(bot.parentId);
    }
    if (parentIds.size === 1) {
      targetParentId = [...parentIds][0];
    } else if (parentIds.size > 1) {
      sendTo(ws, {
        type: "ack",
        requestId: msg.requestId || null,
        ok: false,
        error: "Ambiguous username across multiple parents; specify parentId",
      });
      return;
    }
  }

  if (!targetParentId && parentsById.size === 1) {
    targetParentId = parentsById.keys().next().value;
  }

  const targetWs = targetParentId ? parentsById.get(targetParentId) : null;
  if (!targetWs || targetWs.readyState !== WebSocket.OPEN) {
    sendTo(ws, {
      type: "ack",
      requestId: msg.requestId || null,
      ok: false,
      error: "Target parent not connected",
    });
    return;
  }

  const relayId = nextRelayId(sessionId, msg.requestId || null, targetParentId);
  sendTo(targetWs, { ...msg, parentId: targetParentId, _relayId: relayId });
}

// ─── startup ──────────────────────────────────────────────────────────────────

httpServer.listen(WEB_PORT, () => {
  console.log(`[server] web ui       http://localhost:${WEB_PORT}`);
  console.log(`[server] browser ws   ws://localhost:${WEB_PORT}/ws`);
  console.log(`[server] parent ws    ws://localhost:${WEB_PORT}/parent`);
});

function shutdown() {
  httpServer.close();
  browserWss.close();
  parentWssI.close();
}

process.on("SIGINT", shutdown);
process.on("SIGTERM", shutdown);