@wopr-network/platform-core 1.0.10 → 1.1.0

package/dist/fleet/node-connection-manager.js

@@ -0,0 +1,320 @@
+import { randomUUID } from "node:crypto";
+import { z } from "zod";
+import { logger } from "../config/logger.js";
+import { ConcurrentTransitionError, InvalidTransitionError } from "./node-state-machine.js";
+/** Runtime validation for inbound WebSocket command results */
+// Keep in sync with CommandResult interface above
+const commandResultSchema = z.object({
+    id: z.string(),
+    type: z.literal("command_result"),
+    command: z.string(),
+    success: z.boolean(),
+    data: z.unknown().optional(),
+    error: z.string().optional(),
+});
+/**
+ * Platform-side WebSocket manager for node agents.
+ *
+ * Responsibilities:
+ * - Accept WebSocket connections from node agents
+ * - Receive heartbeat messages and update nodes table
+ * - Send commands to specific node agents and await results
+ * - Track connection state per node
+ */
+/** Statuses that indicate a node has crashed and been recovered by RecoveryManager */
+const CRASHED_STATUSES = new Set(["offline", "recovering", "failed"]);
+export class NodeConnectionManager {
+    nodeRepo;
+    botInstanceRepo;
+    recoveryRepo;
+    connections = new Map();
+    pending = new Map();
+    orphanCleaner;
+    cleanupInFlight = new Set();
+    onNodeRegistered;
+    constructor(nodeRepo, botInstanceRepo, recoveryRepo, options) {
+        this.nodeRepo = nodeRepo;
+        this.botInstanceRepo = botInstanceRepo;
+        this.recoveryRepo = recoveryRepo;
+        this.onNodeRegistered = options?.onNodeRegistered;
+    }
+    /** Inject OrphanCleaner after construction to break the circular dependency. */
+    setOrphanCleaner(cleaner) {
+        this.orphanCleaner = cleaner;
+    }
+    /**
+     * Register a new node (called from HTTP POST /internal/nodes/register)
+     */
+    async registerNode(registration) {
+        const existing = await this.nodeRepo.getById(registration.node_id);
+        if (existing) {
+            const wasInCrashedState = CRASHED_STATUSES.has(existing.status);
+            // Re-register via repo (handles metadata update + state transition)
+            await this.nodeRepo.register({
+                nodeId: registration.node_id,
+                host: registration.host,
+                capacityMb: registration.capacity_mb,
+                agentVersion: registration.agent_version,
+            });
+            // If node was in a dead state, close any in-flight recovery events
+            if (wasInCrashedState) {
+                await this.recoveryRepo.completeInProgressEvents(registration.node_id);
+                logger.info(`Node ${registration.node_id} re-registered as returning (prior: ${existing.status})`);
+            }
+            else {
+                logger.info(`Node ${registration.node_id} re-registered`);
+            }
+        }
+        else {
+            await this.nodeRepo.register({
+                nodeId: registration.node_id,
+                host: registration.host,
+                capacityMb: registration.capacity_mb,
+                agentVersion: registration.agent_version,
+            });
+            logger.info(`Node ${registration.node_id} registered`);
+        }
+        // Trigger auto-retry check for waiting recovery tenants
+        if (this.onNodeRegistered) {
+            this.onNodeRegistered();
+        }
+    }
+    /**
+     * Accept an incoming WebSocket for a node
+     */
+    handleWebSocket(nodeId, ws) {
+        // Close any existing connection for this node
+        const existing = this.connections.get(nodeId);
+        if (existing && existing.readyState === 1 /* OPEN */) {
+            logger.warn(`Closing existing WebSocket for ${nodeId}`);
+            existing.close();
+        }
+        this.connections.set(nodeId, ws);
+        logger.info(`WebSocket connected for node ${nodeId}`);
+        ws.on("message", (data) => {
+            this.handleMessage(nodeId, data);
+        });
+        ws.on("close", () => {
+            logger.info(`WebSocket closed for node ${nodeId}`);
+            // Only remove if this is still the current connection (avoids race with reconnects)
+            if (this.connections.get(nodeId) === ws) {
+                this.connections.delete(nodeId);
+            }
+        });
+        ws.on("error", (err) => {
+            logger.error(`WebSocket error for node ${nodeId}`, { err: err.message });
+        });
+    }
+    /**
+     * Handle an inbound WebSocket message from a node
+     */
+    handleMessage(nodeId, data) {
+        let parsed;
+        try {
+            parsed = JSON.parse(data.toString());
+        }
+        catch {
+            logger.warn(`Received non-JSON message from ${nodeId}`);
+            return;
+        }
+        const msg = parsed;
+        // Handle heartbeat
+        if (msg.type === "heartbeat") {
+            this.processHeartbeat(nodeId, msg).catch((err) => {
+                logger.error(`Heartbeat processing failed for ${nodeId}`, { err });
+            });
+            return;
+        }
+        // Handle command result
+        if (msg.type === "command_result" && typeof msg.id === "string") {
+            const parsed = commandResultSchema.safeParse(msg);
+            if (!parsed.success) {
+                logger.warn(`Malformed command_result from ${nodeId}`, { id: msg.id, issues: parsed.error.issues });
+                return;
+            }
+            const result = parsed.data;
+            const pending = this.pending.get(result.id);
+            if (pending) {
+                clearTimeout(pending.timeout);
+                this.pending.delete(result.id);
+                if (result.success) {
+                    pending.resolve(result);
+                }
+                else {
+                    pending.reject(new Error(result.error ?? "Command failed"));
+                }
+            }
+            return;
+        }
+        // Handle health events (optional logging)
+        if (msg.type === "health_event") {
+            logger.warn(`Health event from ${nodeId}`, { event: msg });
+            return;
+        }
+        logger.debug(`Unknown message type from ${nodeId}`, { msg });
+    }
+    /**
+     * Process a heartbeat message and update nodes table.
+     * For "returning" nodes, triggers OrphanCleaner (fire-and-forget) instead of
+     * immediately setting active — cleanup handles the transition.
+     */
+    async processHeartbeat(nodeId, msg) {
+        const containers = msg.containers;
+        const usedMb = containers?.reduce((sum, c) => sum + (c.memory_mb ?? 0), 0) ?? 0;
+        // Check current node status to decide how to handle the heartbeat
+        const status = await this.nodeRepo.getStatus(nodeId);
+        if (status === null) {
+            logger.warn(`Heartbeat from unknown node ${nodeId}`);
+            return;
+        }
+        // For returning nodes: trigger orphan cleanup (it will transition to active).
+        // Use cleanupInFlight set to prevent double-invocation if two heartbeats race.
+        if (status === "returning" && this.orphanCleaner && !this.cleanupInFlight.has(nodeId)) {
+            this.cleanupInFlight.add(nodeId);
+            const containerNames = (containers ?? []).map((c) => c.name);
+            this.orphanCleaner
+                .clean({ nodeId, runningContainers: containerNames })
+                .catch((err) => {
+                logger.error(`Orphan cleanup failed for node ${nodeId}`, {
+                    err: err instanceof Error ? err.message : String(err),
+                });
+            })
+                .finally(() => {
+                this.cleanupInFlight.delete(nodeId);
+            });
+        }
+        // Transition unhealthy -> active through the state machine (with audit trail).
+        // For "active" nodes: no transition needed, just update heartbeat timestamp.
+        // For "returning" nodes: OrphanCleaner owns the transition, skip here.
+        if (status === "unhealthy") {
+            try {
+                await this.nodeRepo.transition(nodeId, "active", "heartbeat_received", "heartbeat");
+            }
+            catch (err) {
+                if (err instanceof ConcurrentTransitionError || err instanceof InvalidTransitionError) {
+                    // Status changed underneath us — log and continue; the heartbeat timestamp update below still applies.
+                    logger.debug(`Heartbeat transition skipped for ${nodeId}`, {
+                        err: err instanceof Error ? err.message : String(err),
+                    });
+                }
+                else {
+                    throw err;
+                }
+            }
+        }
+        await this.nodeRepo.updateHeartbeat(nodeId, usedMb);
+        logger.debug(`Heartbeat received from ${nodeId}`, { usedMb, status });
+    }
+    /**
+     * Send a command to a node agent and return the result
+     */
+    async sendCommand(nodeId, command, timeoutMs = 60_000) {
+        const ws = this.connections.get(nodeId);
+        if (!ws || ws.readyState !== 1 /* OPEN */) {
+            throw new Error(`Node ${nodeId} is not connected`);
+        }
+        const id = randomUUID();
+        const fullCommand = { id, ...command };
+        return new Promise((resolve, reject) => {
+            const timeout = setTimeout(() => {
+                this.pending.delete(id);
+                reject(new Error(`Command ${command.type} timed out on node ${nodeId}`));
+            }, timeoutMs);
+            this.pending.set(id, { resolve, reject, timeout });
+            ws.send(JSON.stringify(fullCommand));
+            logger.debug(`Sent command ${command.type} to ${nodeId}`, { commandId: id });
+        });
+    }
+    /**
+     * Get current status of all nodes
+     */
+    async listNodes() {
+        const nodes = await this.nodeRepo.list();
+        return nodes;
+    }
+    /**
+     * Get tenants assigned to a specific node
+     */
+    async getNodeTenants(nodeId) {
+        const instances = await this.botInstanceRepo.listByNode(nodeId);
+        return instances.map((inst) => ({
+            id: inst.id,
+            tenantId: inst.tenantId,
+            name: inst.name,
+            containerName: `tenant_${inst.tenantId}`,
+            estimatedMb: 100, // Default estimate; can be refined from heartbeat data
+        }));
+    }
+    /**
+     * Find the best target node for recovery (most free capacity)
+     */
+    async findBestTarget(excludeNodeId, requiredMb) {
+        const node = await this.nodeRepo.findBestTarget(excludeNodeId, requiredMb);
+        return node;
+    }
+    /**
+     * Reassign a tenant to a new node (update bot_instances)
+     */
+    async reassignTenant(botId, targetNodeId) {
+        await this.botInstanceRepo.reassign(botId, targetNodeId);
+        logger.info(`Reassigned bot ${botId} to node ${targetNodeId}`);
+    }
+    /**
+     * Update a node's used capacity
+     */
+    async addNodeCapacity(nodeId, deltaMb) {
+        await this.nodeRepo.addCapacity(nodeId, deltaMb);
+    }
+    /**
+     * Get a single node by ID
+     */
+    async getNode(nodeId) {
+        const node = await this.nodeRepo.getById(nodeId);
+        return node ?? undefined;
+    }
+    /**
+     * Check if a node is connected (has active WebSocket)
+     */
+    isConnected(nodeId) {
+        const ws = this.connections.get(nodeId);
+        return ws != null && ws.readyState === 1;
+    }
+    /**
+     * Register a self-hosted node with owner and per-node secret hash.
+     */
+    async registerSelfHostedNode(registration) {
+        await this.nodeRepo.registerSelfHosted({
+            nodeId: registration.node_id,
+            host: registration.host,
+            capacityMb: registration.capacity_mb,
+            agentVersion: registration.agent_version,
+            ownerUserId: registration.ownerUserId,
+            label: registration.label,
+            nodeSecretHash: registration.nodeSecretHash,
+        });
+        logger.info(`Self-hosted node ${registration.node_id} registered for user ${registration.ownerUserId}`);
+        // Trigger auto-retry check for waiting recovery tenants
+        if (this.onNodeRegistered) {
+            this.onNodeRegistered();
+        }
+    }
+    /**
+     * Look up a node by its persistent secret (hashed). Returns node or undefined.
+     */
+    async getNodeBySecret(secret) {
+        const node = await this.nodeRepo.getBySecret(secret);
+        return node ?? undefined;
+    }
+    /**
+     * Remove a self-hosted node (deregister), closing any active WebSocket.
+     */
+    async removeNode(nodeId) {
+        const ws = this.connections.get(nodeId);
+        if (ws) {
+            ws.close();
+            this.connections.delete(nodeId);
+        }
+        await this.nodeRepo.delete(nodeId);
+        logger.info(`Node ${nodeId} removed`);
+    }
+}

package/dist/fleet/node-connection-manager.test.d.ts

@@ -0,0 +1 @@
+export {};