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 { advanceJobsFrom } from "../../dispatch";
import { enqueueFrameIngestionBatch } from "../../queue";
import { edgeJwtSecret } from "../../auth/edge_secret";

/**
 * 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;
  /**
   * Where this gateway serves camera restreams, from its hello. Lives on the
   * connection rather than in a DB column or a farm-keyed map: it is a LAN
   * address that is only reachable while this gateway is up, so it should
   * vanish with the socket instead of outliving it as a URL nothing answers.
   * The alternative, a column on `farms`, would keep serving a stale address
   * after the gateway disconnected or moved. Undefined means none advertised.
   */
  streamBaseUrl?: string;
  /**
   * Closes this socket if it never authenticates. Cleared by a successful
   * `hello` and on close.
   *
   * The cost of carrying the token in the first message rather than the
   * handshake is exactly this: there is a window in which an open socket
   * belongs to nobody. Every other message type already refuses to act
   * without a `farmId`, so such a socket can do nothing - but it does sit in
   * `connections` forever, and that map is keyed by whoever can reach the
   * port. A silent connection is not a slow one: a real gateway sends
   * `hello` as its first frame.
   */
  handshakeTimer?: ReturnType<typeof setTimeout>;
  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>();

/**
 * farmId -> the printers its gateway is reporting that nothing in the cloud
 * claims, keyed by the gateway's own id for them.
 *
 * `persistTelemetry` deliberately refuses to insert a printer row for an id
 * it does not recognise - a typo in someone's `printers.toml` would otherwise
 * become a permanent phantom printer, and a gateway should not be able to
 * create records just by mentioning them. But the previous behaviour was to
 * log a warning and drop the snapshot, which made the single worst first-run
 * experience in the platform: the gateway connects, reports happily, and the
 * fleet stays empty with the only evidence in the server's stderr.
 *
 * So the refusal stands and the *observation* is kept. A new farm's owner can
 * list what their gateway can see and claim it with `POST /printers`, which
 * is the one place printer rows are created.
 *
 * In memory on purpose: this is live gateway state, exactly like
 * `connections` and `gatewayByFarm` beside it. It is rebuilt within one
 * telemetry interval of a reconnect, and persisting it would mean inventing
 * a second, staler source of truth about what a gateway can see.
 */
const unclaimedByFarm = new Map<string, Map<string, UnclaimedPrinter>>();

export interface UnclaimedPrinter {
  /** The id in the gateway's `printers.toml` - what to put in `gatewayPrinterId`. */
  gatewayPrinterId: string;
  /** The gateway's own display name for it, to make a claim prompt legible. */
  name: string;
  /** Its vendor state, so an owner can tell a live printer from a stale entry. */
  vendorState: string | null;
  firstSeenAt: string;
  lastSeenAt: string;
}

function noteUnclaimed(farmId: string, snapshot: PrinterTelemetry) {
  let farm = unclaimedByFarm.get(farmId);
  if (!farm) {
    farm = new Map();
    unclaimedByFarm.set(farmId, farm);
  }
  const now = new Date().toISOString();
  const existing = farm.get(snapshot.printer_id);
  farm.set(snapshot.printer_id, {
    gatewayPrinterId: snapshot.printer_id,
    name: snapshot.name,
    vendorState: snapshot.state ?? null,
    firstSeenAt: existing?.firstSeenAt ?? now,
    lastSeenAt: now,
  });
}

/**
 * What this farm's gateway is reporting that no printer row claims.
 *
 * Empty is the healthy answer, and also the answer when no gateway has ever
 * connected - a caller that wants to tell those apart should check
 * `isGatewayOnline` too.
 */
export function listUnclaimedPrinters(farmId: string): UnclaimedPrinter[] {
  return [...(unclaimedByFarm.get(farmId)?.values() ?? [])].sort((a, b) =>
    a.gatewayPrinterId.localeCompare(b.gatewayPrinterId),
  );
}

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

/**
 * How long a socket may stay open without authenticating. Generous compared
 * with the round trip it covers - the gateway sends `hello` immediately after
 * the socket opens - because the only thing a tight bound would buy is
 * dropping gateways on slow links.
 */
const HANDSHAKE_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) {
        // Still not inserted - see `unclaimedByFarm` for why - but no longer
        // only a line in stderr. An owner can now see it and claim it.
        noteUnclaimed(conn.farmId, snapshot);
        continue;
      }

      printerId = row.id;
      conn.printerIds.set(snapshot.printer_id, printerId);
      // It is claimed now, so it is no longer waiting to be.
      unclaimedByFarm.get(conn.farmId)?.delete(snapshot.printer_id);
    }

    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,
        // Cleared like the rest when absent: a printer that lost its stream
        // must stop being offered one.
        cameraStream: snapshot.camera_stream ?? null,
        lastSeenAt: new Date(),
        updatedAt: new Date(),
      })
      .where(eq(printers.id, printerId));

    // The printer's row is the latest snapshot; the job's row is the history
    // that outlives it. This is the only thing that moves a job out of
    // `uploading` - see `advanceJobsFrom` for why the gateway's `ok` on a
    // start command is explicitly not trusted to do it.
    await advanceJobsFrom(printerId, status, snapshot.percent);
  }
}

/**
 * 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);
}

/**
 * The stream base URL the farm's currently connected gateway advertised, or
 * undefined if no gateway is connected or it advertised none. Deliberately
 * does not distinguish those two: callers only need to know it is unusable.
 */
export function getStreamBaseUrl(farmId: string): string | undefined {
  const wsId = gatewayByFarm.get(farmId);
  return wsId ? connections.get(wsId)?.streamBaseUrl : undefined;
}

export const edgeWsRoute = new Elysia()
  .use(jwt({ name: "edgeJwt", secret: edgeJwtSecret }))
  .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.
      const conn: EdgeConn = {
        farmId: "",
        gatewayId: "",
        sessionId: crypto.randomUUID(),
        lastSeenAt: Date.now(),
        send: (message) => ws.send(message),
        printerIds: new Map(),
        pending: new Map(),
      };

      conn.handshakeTimer = setTimeout(() => {
        if (conn.farmId) return;
        console.warn(`[edge] closing socket ${ws.id}: no hello within ${HANDSHAKE_TIMEOUT_MS}ms`);
        ws.close();
      }, HANDSHAKE_TIMEOUT_MS);

      connections.set(ws.id, conn);
    },

    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;
          }

          // Authenticated, so it is no longer a socket belonging to nobody.
          // Cleared only here, after the token verified and the farm matched
          // - a rejected hello leaves the timer running, so a client that
          // retries a bad token on one socket still gets dropped.
          clearTimeout(conn.handshakeTimer);
          conn.handshakeTimer = undefined;

          conn.farmId = farmId;
          conn.gatewayId = message.gateway_id;
          // `?? undefined`: absent and explicit null both mean "no streams".
          conn.streamBaseUrl = message.stream_base_url ?? undefined;

          // 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);
          // `?? undefined` because the wire allows either absent or null for
          // "no error" (see the Maybe() helper in @continuum/ts-types) and
          // callers should only have to check one of them.
          pending.resolve({ ok: message.ok, error: message.error ?? undefined });
          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;
          }

          // One round trip for the whole batch, and snake_case field names:
          // this is the `stream:camera_frames` contract in
          // `schemas/continuum/v1/ai/`, not an HTTP response, and the
          // consumer at the other end is serde-derived Rust.
          await enqueueFrameIngestionBatch(
            message.frames.map((frame) => ({
              farm_id: conn.farmId,
              printer_id: printerId,
              // Omitted rather than null when there is no job: a stream
              // entry has no null, and `""` is a value a consumer would
              // believe. `?? undefined` collapses the wire's "absent or
              // null" into the one the schema describes.
              ...(message.print_job_id != null ? { print_job_id: message.print_job_id } : {}),
              r2_object_key: frame.object_key,
              captured_at: frame.captured_at,
            })),
          );
          break;
        }
      }
    },

    close(ws) {
      const conn = connections.get(ws.id);
      if (conn) {
        clearTimeout(conn.handshakeTimer);
        // 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);
    },
  });