continuum-backend/src/routes/ws/edge.ts

import { Elysia } from "elysia";
import jwt from "@elysiajs/jwt";
import { and, eq } from "drizzle-orm";
import {
  EdgeClientMessageSchema,
  type EdgeServerMessage,
  type PrinterCommand,
  type PrinterTelemetry,
} from "@continuum/ts-types";
import { db } from "../../db";
import { printers } from "../../db/schema";
import { enqueueFrameIngestion } from "../../queue";

/**
 * Edge-gateway uplink. A continuum-proxy gateway connects here with a
 * short-lived edge JWT (minted via POST /auth/edge-token), streams telemetry
 * for its whole printer fleet, and receives commands for individual printers
 * on the same socket.
 *
 * The wire format is `schemas/continuum/v1/edge/` in continuum-schemas,
 * validated here through its TypeBox mirror in `@continuum/ts-types`. It is
 * snake_case, unlike this backend's HTTP routes: the other end is
 * serde-derived Rust, and this route previously drifted into a protocol the
 * gateway did not speak at all (an `auth` handshake it never sends, one
 * printer per telemetry message instead of a batch, UUID printer ids a
 * gateway has no way to know, and no command path whatsoever) precisely
 * because there was nothing written down for both sides to mirror.
 */

interface PendingCommand {
  resolve: (result: { ok: boolean; error?: string }) => void;
  timer: ReturnType<typeof setTimeout>;
}

interface EdgeConn {
  farmId: string;
  gatewayId: string;
  sessionId: string;
  lastSeenAt: number;
  send: (message: EdgeServerMessage) => void;
  /**
   * Gateway printer id -> our `printers.id`. Filled on the first telemetry
   * batch and reused after that, so a steady-state connection does not
   * re-query the same mapping every few seconds.
   */
  printerIds: Map<string, string>;
  /** Commands sent but not yet answered, keyed by command_id. */
  pending: Map<string, PendingCommand>;
}

const connections = new Map<string, EdgeConn>();

/**
 * farmId -> ws.id. One gateway per farm, so a second connection for the same
 * farm replaces the first: that is what a gateway reconnecting after a
 * network drop looks like, and the stale entry would otherwise keep receiving
 * commands nothing is listening for.
 */
const gatewayByFarm = new Map<string, string>();

/** How long to wait for a gateway to answer a command before giving up. */
const COMMAND_TIMEOUT_MS = 10_000;

/**
 * Maps a printer's own state string onto the narrowed `printer_status` enum.
 *
 * Unrecognised states deliberately fall through to the current value rather
 * than to "idle": claiming a printer is idle when we do not understand what
 * it said is the one wrong answer here, since idle is what the scheduler
 * treats as "safe to send work to". The raw string is kept in
 * `vendor_state` either way.
 */
function mapStatus(state: string | null | undefined, printError: number | null | undefined) {
  // `!= null` deliberately, not `!== undefined`: these fields arrive as either
  // absent or an explicit null (see the Maybe() helper in @continuum/ts-types),
  // and treating a null error code as a non-zero one would report every
  // printer that has nothing to say as faulted.
  if (printError != null && printError !== 0) return "error" as const;

  switch (state?.toUpperCase()) {
    case "IDLE":
    case "FINISH":
      return "idle" as const;
    case "RUNNING":
    case "PREPARE":
    case "SLICING":
      return "printing" as const;
    case "PAUSE":
      return "paused" as const;
    case "FAILED":
      return "error" as const;
    default:
      return undefined;
  }
}

/**
 * Writes one telemetry batch to the `printers` table.
 *
 * Only rows whose `gateway_printer_id` is already claimed by this farm are
 * updated - an unknown id is logged and skipped rather than inserted. A
 * gateway should not be able to create printer records just by mentioning
 * them: that is what POST /printers is for, and auto-inserting here would
 * turn a typo in someone's printers.toml into a permanent phantom printer.
 */
async function persistTelemetry(conn: EdgeConn, batch: PrinterTelemetry[]) {
  for (const snapshot of batch) {
    let printerId = conn.printerIds.get(snapshot.printer_id);

    if (!printerId) {
      const row = await db.query.printers.findFirst({
        columns: { id: true },
        where: and(
          eq(printers.farmId, conn.farmId),
          eq(printers.gatewayPrinterId, snapshot.printer_id),
        ),
      });

      if (!row) {
        console.warn(
          `[edge] gateway ${conn.gatewayId} reported unknown printer ` +
            `'${snapshot.printer_id}'; no printer in farm ${conn.farmId} claims that ` +
            `gateway_printer_id, skipping`,
        );
        continue;
      }

      printerId = row.id;
      conn.printerIds.set(snapshot.printer_id, printerId);
    }

    const status = mapStatus(snapshot.state, snapshot.print_error);

    await db
      .update(printers)
      .set({
        // `?? null` rather than leaving the key out: a field the gateway
        // stopped reporting has genuinely gone away (a finished job has no
        // remaining minutes), and keeping the last value would leave stale
        // numbers on screen forever.
        ...(status ? { status } : {}),
        vendorState: snapshot.state ?? null,
        vendorErrorCode: snapshot.print_error ?? null,
        nozzleTempC: snapshot.nozzle_temp ?? null,
        nozzleTargetTempC: snapshot.nozzle_target ?? null,
        bedTempC: snapshot.bed_temp ?? null,
        bedTargetTempC: snapshot.bed_target ?? null,
        progressPercent: snapshot.percent ?? null,
        currentLayer: snapshot.layer ?? null,
        totalLayers: snapshot.total_layers ?? null,
        remainingMinutes: snapshot.remaining_minutes ?? null,
        currentJobName: snapshot.job_name ?? null,
        activeAlerts: snapshot.active_alerts ?? null,
        lastSeenAt: new Date(),
        updatedAt: new Date(),
      })
      .where(eq(printers.id, printerId));
  }
}

/**
 * Sends a command to the printer's gateway and waits for its result.
 *
 * Takes our own `printers.id` and translates to the gateway's id on the way
 * out, so no caller elsewhere in the backend has to know that gateways use
 * different identifiers.
 *
 * Throws if the printer is unknown, unclaimed, or its farm has no gateway
 * connected right now. Resolves `{ ok: false }` if the gateway answered with
 * a failure or did not answer within `COMMAND_TIMEOUT_MS` - a command that
 * was delivered but never acknowledged is not the same as one that was never
 * sent, and callers need to be able to tell those apart.
 */
export async function sendPrinterCommand(
  printerId: string,
  command: PrinterCommand,
): Promise<{ ok: boolean; error?: string; commandId: string }> {
  const printer = await db.query.printers.findFirst({
    columns: { id: true, farmId: true, gatewayPrinterId: true },
    where: eq(printers.id, printerId),
  });

  if (!printer) throw new Error(`printer ${printerId} not found`);
  if (!printer.gatewayPrinterId) {
    throw new Error(`printer ${printerId} has no gateway_printer_id - no gateway has claimed it`);
  }

  const wsId = gatewayByFarm.get(printer.farmId);
  const conn = wsId ? connections.get(wsId) : undefined;
  if (!conn) throw new Error(`no gateway connected for farm ${printer.farmId}`);

  const commandId = crypto.randomUUID();

  const result = await new Promise<{ ok: boolean; error?: string }>((resolve) => {
    const timer = setTimeout(() => {
      conn.pending.delete(commandId);
      resolve({ ok: false, error: `gateway did not answer within ${COMMAND_TIMEOUT_MS}ms` });
    }, COMMAND_TIMEOUT_MS);

    conn.pending.set(commandId, { resolve, timer });

    conn.send({
      type: "command",
      command_id: commandId,
      printer_id: printer.gatewayPrinterId!,
      command,
    });
  });

  return { ...result, commandId };
}

/** Whether a farm currently has a gateway connected. */
export function isGatewayOnline(farmId: string): boolean {
  const wsId = gatewayByFarm.get(farmId);
  return wsId !== undefined && connections.has(wsId);
}

export const edgeWsRoute = new Elysia()
  .use(jwt({ name: "edgeJwt", secret: process.env.EDGE_JWT_SECRET ?? "dev-edge-secret-change-me" }))
  .ws("/ws/edge/v1", {
    body: EdgeClientMessageSchema,

    open(ws) {
      // No farmId yet - set by `hello`, and every other message type refuses
      // to do anything until it is.
      connections.set(ws.id, {
        farmId: "",
        gatewayId: "",
        sessionId: crypto.randomUUID(),
        lastSeenAt: Date.now(),
        send: (message) => ws.send(message),
        printerIds: new Map(),
        pending: new Map(),
      });
    },

    async message(ws, message) {
      const conn = connections.get(ws.id);
      if (!conn) return;

      conn.lastSeenAt = Date.now();

      switch (message.type) {
        case "hello": {
          const payload = await ws.data.edgeJwt.verify(message.token);
          if (!payload || payload.scope !== "edge") {
            ws.send({ type: "hello_error", message: "invalid edge token" });
            ws.close();
            return;
          }

          // The token decides the farm, never the message body: a gateway
          // claiming a farm_id its token was not minted for is exactly what
          // this check is here to stop.
          const farmId = payload.farmId as string;
          if (message.farm_id && message.farm_id !== farmId) {
            ws.send({ type: "hello_error", message: "farm_id does not match the token's farm" });
            ws.close();
            return;
          }

          conn.farmId = farmId;
          conn.gatewayId = message.gateway_id;

          // One gateway per farm: drop whatever was registered before. On a
          // reconnect that stale socket is already dead, and leaving it
          // registered would send commands into it.
          const previous = gatewayByFarm.get(farmId);
          if (previous && previous !== ws.id) {
            console.warn(
              `[edge] farm ${farmId} already had gateway connection ${previous}; ` +
                `replacing it with ${ws.id} (gateway ${message.gateway_id})`,
            );
          }
          gatewayByFarm.set(farmId, ws.id);

          ws.send({ type: "hello_ack", session_id: conn.sessionId });
          console.info(`[edge] gateway ${message.gateway_id} connected for farm ${farmId}`);
          break;
        }

        case "heartbeat": {
          // Echo the sequence number back. The gateway uses it to tell a
          // stalled connection from a slow one, which it cannot do against a
          // bare ack.
          ws.send({ type: "heartbeat_ack", seq: message.seq });
          break;
        }

        case "telemetry": {
          if (!conn.farmId) return;
          await persistTelemetry(conn, message.printers);
          break;
        }

        case "command_result": {
          const pending = conn.pending.get(message.command_id);
          if (!pending) {
            // Either the command already timed out, or this gateway answered
            // something we never asked. Both are worth saying out loud.
            console.warn(
              `[edge] gateway ${conn.gatewayId} sent a result for unknown or ` +
                `expired command ${message.command_id}`,
            );
            return;
          }
          clearTimeout(pending.timer);
          conn.pending.delete(message.command_id);
          pending.resolve({ ok: message.ok, error: message.error });
          break;
        }

        case "frame_batch": {
          if (!conn.farmId) return;

          const printerId = conn.printerIds.get(message.printer_id);
          if (!printerId) {
            console.warn(
              `[edge] frame_batch for unknown printer '${message.printer_id}' from ` +
                `gateway ${conn.gatewayId}, dropping ${message.frames.length} frame(s)`,
            );
            return;
          }

          for (const frame of message.frames) {
            await enqueueFrameIngestion({
              farmId: conn.farmId,
              printerId,
              printJobId: message.print_job_id ?? null,
              r2ObjectKey: frame.object_key,
              capturedAt: frame.captured_at,
            });
          }
          break;
        }
      }
    },

    close(ws) {
      const conn = connections.get(ws.id);
      if (conn) {
        // Anything still waiting on this gateway will never be answered now.
        // Failing them immediately beats making every caller wait out the
        // full timeout for a socket that is already gone.
        for (const [commandId, pending] of conn.pending) {
          clearTimeout(pending.timer);
          pending.resolve({ ok: false, error: "gateway disconnected before answering" });
          conn.pending.delete(commandId);
        }
        if (conn.farmId && gatewayByFarm.get(conn.farmId) === ws.id) {
          gatewayByFarm.delete(conn.farmId);
        }
        if (conn.gatewayId) {
          console.info(`[edge] gateway ${conn.gatewayId} disconnected (farm ${conn.farmId})`);
        }
      }
      connections.delete(ws.id);
    },
  });