pikasTech-HWLAB/internal/cloud/hwpod-node-ws-registry.ts

import { randomUUID } from "node:crypto";

const DEFAULT_DISPATCH_TIMEOUT_MS = 30000;

type HwpodNodeConnection = {
  id: string;
  socket: any;
  nodeId: string | null;
  name: string | null;
  capabilities: string[];
  labels: Record<string, unknown>;
  firstSeenAt: string;
  lastSeenAt: string;
};

type PendingDispatch = {
  nodeId: string;
  requestId: string;
  timer: ReturnType<typeof setTimeout>;
  resolve: (value: any) => void;
};

export function createHwpodNodeWsRegistry(options: any = {}) {
  const now = options.now ?? (() => new Date().toISOString());
  const clock = options.clock ?? (() => Date.now());
  const connections = new Map<string, HwpodNodeConnection>();
  const pending = new Map<string, PendingDispatch>();
  const socketConnections = new WeakMap<object, HwpodNodeConnection>();

  function openBunSocket(socket: any) {
    const connection: HwpodNodeConnection = {
      id: `hwpod_node_conn_${randomUUID()}`,
      socket,
      nodeId: null,
      name: null,
      capabilities: [],
      labels: {},
      firstSeenAt: now(),
      lastSeenAt: now()
    };
    socketConnections.set(socket, connection);
    return connection;
  }

  function handleBunMessage(socket: any, raw: unknown) {
    const connection = socketConnections.get(socket);
    if (!connection) return;
    const text = typeof raw === "string" ? raw : Buffer.from(raw as any).toString("utf8");
    handleText(connection, text);
  }

  function closeBunSocket(socket: any) {
    const connection = socketConnections.get(socket);
    if (!connection) return;
    socketConnections.delete(socket);
    removeConnection(connection);
  }

  function dispatch(plan: any, requestMeta: any = {}, dispatchOptions: any = {}) {
    const nodeId = safeNodeId(plan?.nodeId);
    const connection = nodeId ? connections.get(nodeId) : null;
    if (!nodeId || !connection) {
      return Promise.resolve(blockedDispatch(plan, requestMeta, `hwpod-node ${nodeId || "<missing>"} is not connected through outbound WebSocket`));
    }
    const requestId = safeText(requestMeta.requestId) || `req_hwpod_ws_${randomUUID()}`;
    const timeoutMs = positiveInteger(dispatchOptions.timeoutMs, DEFAULT_DISPATCH_TIMEOUT_MS);
    return new Promise((resolve) => {
      const timer = setTimeout(() => {
        pending.delete(requestId);
        resolve(blockedDispatch(plan, requestMeta, `hwpod-node ${nodeId} WebSocket dispatch timed out after ${timeoutMs}ms`));
      }, timeoutMs);
      pending.set(requestId, { nodeId, requestId, timer, resolve });
      try {
        sendJson(connection, { type: "hwpod-node-ops", requestId, plan, requestMeta });
      } catch (error) {
        clearTimeout(timer);
        pending.delete(requestId);
        resolve(blockedDispatch(plan, requestMeta, error instanceof Error ? error.message : String(error)));
      }
    });
  }

  function describe() {
    const observedAt = now();
    return {
      mode: "hwpod-node-outbound-native-ws",
      route: "/v1/hwpod-node/ws",
      connectedCount: connections.size,
      pendingCount: pending.size,
      nodes: [...connections.values()].map((connection) => ({
        nodeId: connection.nodeId,
        name: connection.name,
        status: "online",
        capabilityCount: connection.capabilities.length,
        capabilities: connection.capabilities,
        labels: connection.labels,
        firstSeenAt: connection.firstSeenAt,
        lastSeenAt: connection.lastSeenAt,
        ageMs: Math.max(0, clock() - Date.parse(connection.lastSeenAt || observedAt))
      }))
    };
  }

  function hasNode(nodeId: string) {
    return connections.has(nodeId);
  }

  function handleText(connection: HwpodNodeConnection, text: string) {
    let message: any;
    try {
      message = JSON.parse(text);
    } catch (error) {
      sendJson(connection, { type: "ack", requestId: "message", ok: false, message: "invalid JSON message" });
      return;
    }
    const messageType = safeText(message.type);
    const nodeId = safeNodeId(message.nodeId);
    if (messageType === "register") {
      if (!nodeId) {
        sendJson(connection, { type: "ack", requestId: "register", ok: false, message: "nodeId is required" });
        return;
      }
      connection.nodeId = nodeId;
      connection.name = safeText(message.name) || nodeId;
      connection.capabilities = Array.isArray(message.capabilities) ? message.capabilities.map(String).filter(Boolean).slice(0, 64) : [];
      connection.labels = message.labels && typeof message.labels === "object" && !Array.isArray(message.labels) ? message.labels : {};
      connection.lastSeenAt = now();
      const previous = connections.get(nodeId);
      if (previous && previous !== connection) previous.socket.close();
      connections.set(nodeId, connection);
      sendJson(connection, { type: "ack", requestId: "register", ok: true, message: "registered", nodeId });
      return;
    }
    if (messageType === "heartbeat") {
      if (nodeId && connection.nodeId === nodeId) connection.lastSeenAt = now();
      return;
    }
    if (messageType === "hwpod-node-ops-result") completeDispatch(message);
  }

  function completeDispatch(message: any) {
    const requestId = safeText(message.requestId);
    const waiter = requestId ? pending.get(requestId) : null;
    if (!waiter) return;
    clearTimeout(waiter.timer);
    pending.delete(requestId);
    const result = message.result && typeof message.result === "object" && !Array.isArray(message.result)
      ? message.result
      : { ok: false, status: "failed", results: [], blocker: { code: "hwpod_node_result_invalid", layer: "hwpod-node", retryable: true, summary: "hwpod-node WebSocket result payload is invalid" } };
    waiter.resolve(result);
  }

  function removeConnection(connection: HwpodNodeConnection) {
    if (connection.nodeId && connections.get(connection.nodeId) === connection) connections.delete(connection.nodeId);
    for (const [requestId, waiter] of pending) {
      if (waiter.nodeId !== connection.nodeId) continue;
      clearTimeout(waiter.timer);
      pending.delete(requestId);
      waiter.resolve(blockedDispatch({ nodeId: waiter.nodeId, ops: [] }, { requestId }, `hwpod-node ${waiter.nodeId} WebSocket disconnected before returning result`));
    }
  }

  return { openBunSocket, handleBunMessage, closeBunSocket, dispatch, describe, hasNode };
}

function sendJson(connection: HwpodNodeConnection, value: any) {
  connection.socket.send(JSON.stringify(value));
}

function blockedDispatch(plan: any, requestMeta: any, summary: string) {
  const ops = Array.isArray(plan?.ops) ? plan.ops : [];
  return {
    ok: false,
    status: "blocked",
    results: ops.map((op: any, index: number) => ({
      opId: safeText(op?.opId) || `op_${index + 1}`,
      op: safeText(op?.op) || "unknown",
      ok: false,
      status: "blocked",
      blocker: { code: "hwpod_node_unavailable", layer: "hwpod-node", retryable: true, summary }
    })),
    blocker: { code: "hwpod_node_unavailable", layer: "hwpod-node", retryable: true, summary },
    requestMeta
  };
}

function safeText(value: unknown) {
  return typeof value === "string" ? value.trim() : "";
}

function safeNodeId(value: unknown) {
  const text = safeText(value);
  return /^[A-Za-z0-9._:-]{1,128}$/u.test(text) ? text : "";
}

function positiveInteger(value: unknown, fallback: number) {
  const numberValue = Number(value);
  return Number.isInteger(numberValue) && numberValue > 0 ? numberValue : fallback;
}