import type { WebSocket } from "ws"; import type { ControlMessage, ControlMessageType } from "@media-cluster/protocol"; import { createMessage } from "@media-cluster/protocol"; import type { CreatePlaybackSessionPayload } from "@media-cluster/shared-types"; import { prisma } from "../database/client"; import { authenticateNode } from "../nodes/routes"; import { LibrarySyncService } from "../media/sync"; import { MetadataService } from "../metadata/service"; import type { Config } from "../config"; import type { FullLibrarySyncPayload, HeartbeatPayload, HelloPayload, LibraryEventPayload, PlaybackSessionEndedPayload, CreatePlaybackSessionAckPayload, } from "@media-cluster/protocol"; interface NodeConnection { nodeId: string; socket: WebSocket; lastHeartbeat: Date; } interface PendingAck { resolve: (ok: boolean) => void; timer: ReturnType; } interface SyncState { syncId: string; fileIds: Set; } class NodeConnectionManager { private connections = new Map(); private librarySync: LibrarySyncService | null = null; private config: Config | null = null; private offlineCheckInterval: ReturnType | null = null; private pendingSessionAcks = new Map(); private syncState = new Map(); private syncQueues = new Map>(); init(config: Config): void { this.config = config; const metadata = new MetadataService(config.TMDB_API_KEY); this.librarySync = new LibrarySyncService(metadata); this.offlineCheckInterval = setInterval(() => { void this.checkOfflineNodes(); }, 30_000); } async handleConnection(socket: WebSocket): Promise { let nodeId: string | null = null; let authenticated = false; const authTimeout = setTimeout(() => { if (!authenticated) socket.close(4001, "Authentication timeout"); }, 30_000); socket.on("message", async (data) => { try { const message = JSON.parse(data.toString()) as ControlMessage; await this.handleMessage(socket, message, (id) => { nodeId = id; authenticated = true; clearTimeout(authTimeout); }); } catch (err) { console.error("WebSocket message error:", err); this.send(socket, createMessage("ERROR", { message: "Invalid message" })); } }); socket.on("close", () => { if (nodeId) { this.connections.delete(nodeId); this.syncState.delete(nodeId); void prisma.node.update({ where: { id: nodeId }, data: { status: "OFFLINE" }, }); void prisma.mediaFile.updateMany({ where: { nodeId }, data: { available: false }, }); console.log(`Node disconnected: ${nodeId}`); } }); } private async handleMessage( socket: WebSocket, message: ControlMessage, onAuth: (nodeId: string) => void ): Promise { if (message.protocolVersion !== 1) { this.send( socket, createMessage("ERROR", { code: "UNSUPPORTED_PROTOCOL", message: `Unsupported protocol version: ${message.protocolVersion}`, }) ); socket.close(4002, "Unsupported protocol"); return; } switch (message.type) { case "HELLO": await this.handleHello(socket, message.payload as HelloPayload, onAuth); break; case "HEARTBEAT": await this.handleHeartbeat(message.payload as HeartbeatPayload); this.send(socket, createMessage("HEARTBEAT_ACK", {})); break; case "LIBRARY_EVENT": await this.handleLibraryEvent( message.payload as LibraryEventPayload, socket ); break; case "FULL_LIBRARY_SYNC": await this.handleFullSync(message.payload as FullLibrarySyncPayload, socket); break; case "PLAYBACK_SESSION_ENDED": await this.handlePlaybackEnded(message.payload as PlaybackSessionEndedPayload); break; case "CREATE_PLAYBACK_SESSION_ACK": this.handlePlaybackSessionAck(message.payload as CreatePlaybackSessionAckPayload); break; default: this.send(socket, createMessage("ERROR", { message: `Unknown type: ${message.type}` })); } } private async handleHello( socket: WebSocket, payload: HelloPayload, onAuth: (nodeId: string) => void ): Promise { const valid = await authenticateNode(payload.nodeId, payload.apiKey); if (!valid) { this.send( socket, createMessage("HELLO_ACK", { accepted: false, reason: "Invalid credentials", heartbeatIntervalSeconds: 30, fullScanIntervalSeconds: 0, fullScanInterval: "off", fullScanAt: "03:30", }) ); socket.close(4003, "Authentication failed"); return; } this.syncState.delete(payload.nodeId); this.connections.set(payload.nodeId, { nodeId: payload.nodeId, socket, lastHeartbeat: new Date(), }); onAuth(payload.nodeId); await prisma.node.update({ where: { id: payload.nodeId }, data: { status: "ONLINE", lastSeen: new Date(), version: payload.softwareVersion, architecture: payload.architecture, hostname: payload.hostname, // Do not overwrite publicStreamUrl on HELLO — admin may correct typos // (e.g. node01 vs node1) without recreating the Unraid container env. }, }); // Restore availability for previously indexed files until full sync arrives await prisma.mediaFile.updateMany({ where: { nodeId: payload.nodeId }, data: { available: true }, }); console.log(`Node connected: ${payload.nodeId} (${payload.hostname})`); let node = await prisma.node.findUnique({ where: { id: payload.nodeId }, include: { scanRoots: { include: { shelf: true } } }, }); const { ensureDefaultShelves, DEFAULT_MOVIE_SHELF_ID, DEFAULT_SERIES_SHELF_ID, getNodeScanRootsPayload, deriveLegacyPaths } = await import("../library/shelves"); await ensureDefaultShelves(); // Seed master path lists from node env/config when admin has not set them yet. if ( node && (node.scanRoots?.length ?? 0) === 0 && (node.moviesPaths?.length ?? 0) === 0 && (node.seriesPaths?.length ?? 0) === 0 && ((payload.moviesPaths?.length ?? 0) > 0 || (payload.seriesPaths?.length ?? 0) > 0) ) { const rows: Array<{ nodeId: string; path: string; shelfId: string }> = []; for (const p of payload.moviesPaths ?? []) { rows.push({ nodeId: payload.nodeId, path: p, shelfId: DEFAULT_MOVIE_SHELF_ID }); } for (const p of payload.seriesPaths ?? []) { rows.push({ nodeId: payload.nodeId, path: p, shelfId: DEFAULT_SERIES_SHELF_ID }); } if (rows.length) { await prisma.nodeScanRoot.createMany({ data: rows }); const legacy = deriveLegacyPaths( rows.map((r) => ({ path: r.path, kind: r.shelfId === DEFAULT_SERIES_SHELF_ID ? ("series" as const) : ("movie" as const), })) ); node = await prisma.node.update({ where: { id: payload.nodeId }, data: { moviesPaths: legacy.moviesPaths, seriesPaths: legacy.seriesPaths }, include: { scanRoots: { include: { shelf: true } } }, }); } } else if (node && (node.scanRoots?.length ?? 0) === 0 && ((node.moviesPaths?.length ?? 0) > 0 || (node.seriesPaths?.length ?? 0) > 0)) { // Migrate legacy arrays → scan roots once const rows: Array<{ nodeId: string; path: string; shelfId: string }> = []; for (const p of node.moviesPaths) { rows.push({ nodeId: node.id, path: p, shelfId: DEFAULT_MOVIE_SHELF_ID }); } for (const p of node.seriesPaths) { rows.push({ nodeId: node.id, path: p, shelfId: DEFAULT_SERIES_SHELF_ID }); } if (rows.length) { await prisma.nodeScanRoot.createMany({ data: rows }); node = await prisma.node.findUnique({ where: { id: payload.nodeId }, include: { scanRoots: { include: { shelf: true } } }, }); } } const fullScanInterval = node?.fullScanInterval || "off"; const fullScanAt = node?.fullScanAt || "03:30"; const legacySeconds = fullScanInterval === "off" ? 0 : (() => { const m = fullScanInterval.match(/^(\d+)h$/i); return m ? parseInt(m[1], 10) * 3600 : 0; })(); const rootsPayload = await getNodeScanRootsPayload(payload.nodeId); const legacy = deriveLegacyPaths(rootsPayload); this.send( socket, createMessage("HELLO_ACK", { accepted: true, heartbeatIntervalSeconds: this.config?.NODE_HEARTBEAT_INTERVAL_SECONDS ?? 30, fullScanIntervalSeconds: legacySeconds, fullScanInterval, fullScanAt, moviesPaths: legacy.moviesPaths, seriesPaths: legacy.seriesPaths, scanRoots: rootsPayload.map((r) => ({ path: r.path, shelfId: r.shelfId, kind: r.kind, })), }) ); } private async handleHeartbeat(payload: HeartbeatPayload): Promise { const conn = this.connections.get(payload.nodeId); if (conn) conn.lastHeartbeat = new Date(); await prisma.node.update({ where: { id: payload.nodeId }, data: { status: "ONLINE", lastSeen: new Date(), totalStorage: BigInt(payload.storageTotalBytes), freeStorage: BigInt(payload.storageFreeBytes), activeStreams: payload.activeStreams, currentBandwidth: BigInt(payload.currentBytesPerSec), libraryFileCount: payload.libraryFileCount, version: payload.softwareVersion, }, }); } private async handleLibraryEvent( payload: LibraryEventPayload, socket: WebSocket ): Promise { const conn = [...this.connections.values()].find((c) => c.socket === socket); if (!conn || !this.librarySync) return; // Don't request another full sync while one is already in progress if (this.syncState.has(conn.nodeId)) { await this.librarySync.processEvents(conn.nodeId, payload.events); this.send(socket, createMessage("LIBRARY_EVENT_ACK", { accepted: true })); return; } const node = await prisma.node.findUnique({ where: { id: conn.nodeId } }); if (!node) return; const expectedRevision = node.lastRevision + 1; const minRevision = Math.min(...payload.events.map((e) => e.nodeRevision)); if (payload.events.length > 0 && minRevision > expectedRevision) { this.send(socket, createMessage("FULL_LIBRARY_SYNC", { full: true })); return; } await this.librarySync.processEvents(conn.nodeId, payload.events); this.send(socket, createMessage("LIBRARY_EVENT_ACK", { accepted: true })); } private async handleFullSync( payload: FullLibrarySyncPayload, socket: WebSocket ): Promise { const conn = [...this.connections.values()].find((c) => c.socket === socket); if (!conn || !this.librarySync) return; const syncId = payload.syncId ?? `legacy-${Date.now()}`; const isLast = payload.isLast ?? true; const batchIndex = payload.batchIndex ?? 0; let state = this.syncState.get(conn.nodeId); if (!state || state.syncId !== syncId) { state = { syncId, fileIds: new Set() }; this.syncState.set(conn.nodeId, state); } // ACK immediately so the node WS stays free for playback / next batches. // Metadata matching can take seconds per batch and must not block the socket. this.send( socket, createMessage("FULL_LIBRARY_SYNC_ACK", { accepted: true, syncId: payload.syncId, batchIndex: payload.batchIndex, }) ); const prev = this.syncQueues.get(conn.nodeId) ?? Promise.resolve(); const job = prev .then(async () => { if (!this.librarySync) return; await this.librarySync.processFullSync( conn.nodeId, payload.nodeRevision, payload.files, { isLast, fileIdsSoFar: state!.fileIds } ); if (isLast) { console.log( `Full library sync complete for ${conn.nodeId}: ${state!.fileIds.size} files` + (payload.batchCount != null ? ` (${payload.batchCount} batches)` : "") ); this.syncState.delete(conn.nodeId); } else if (batchIndex % 10 === 0) { console.log( `Full library sync progress for ${conn.nodeId}: batch ${(batchIndex ?? 0) + 1}/${payload.batchCount ?? "?"}` ); } }) .catch((err) => { console.error(`Full sync failed for ${conn.nodeId}:`, err); this.syncState.delete(conn.nodeId); }); this.syncQueues.set(conn.nodeId, job); } private handlePlaybackSessionAck(payload: CreatePlaybackSessionAckPayload): void { const pending = this.pendingSessionAcks.get(payload.sessionId); if (!pending) return; clearTimeout(pending.timer); this.pendingSessionAcks.delete(payload.sessionId); pending.resolve(payload.ok); } private async handlePlaybackEnded(payload: PlaybackSessionEndedPayload): Promise { await prisma.playbackSession.updateMany({ where: { id: payload.sessionId }, data: { status: "COMPLETED" }, }); } sendToNode(nodeId: string, type: ControlMessageType, payload: unknown): boolean { const conn = this.connections.get(nodeId); if (!conn || conn.socket.readyState !== 1) return false; this.send(conn.socket, createMessage(type, payload)); return true; } sendRescan(nodeId: string, full: boolean): boolean { return this.sendToNode(nodeId, "RESCAN", { full }); } sendConfigUpdate( nodeId: string, payload: { fullScanInterval?: string; fullScanAt?: string; moviesPaths?: string[]; seriesPaths?: string[]; scanRoots?: Array<{ path: string; shelfId: string; kind: "movie" | "series" }>; } ): boolean { return this.sendToNode(nodeId, "CONFIG_UPDATE", payload); } sendRestart(nodeId: string, reason = "admin"): boolean { return this.sendToNode(nodeId, "RESTART", { reason }); } createPlaybackSession(nodeId: string, payload: CreatePlaybackSessionPayload): boolean { return this.sendToNode(nodeId, "CREATE_PLAYBACK_SESSION", payload); } /** Push session to node and wait for ACK (or timeout). */ async createPlaybackSessionAcked( nodeId: string, payload: CreatePlaybackSessionPayload, timeoutMs = 5000 ): Promise { if (!this.isOnline(nodeId)) return false; const acked = new Promise((resolve) => { const timer = setTimeout(() => { // Fail closed: never hand Stremio a node URL without a confirmed session. this.pendingSessionAcks.delete(payload.sessionId); resolve(false); }, timeoutMs); this.pendingSessionAcks.set(payload.sessionId, { resolve, timer }); }); const pushed = this.createPlaybackSession(nodeId, payload); if (!pushed) { const pending = this.pendingSessionAcks.get(payload.sessionId); if (pending) { clearTimeout(pending.timer); this.pendingSessionAcks.delete(payload.sessionId); } return false; } return acked; } revokePlaybackSession(nodeId: string, sessionId: string): boolean { return this.sendToNode(nodeId, "REVOKE_PLAYBACK_SESSION", { sessionId }); } disconnect(nodeId: string): void { const conn = this.connections.get(nodeId); if (conn) { conn.socket.close(4004, "Revoked"); this.connections.delete(nodeId); this.syncState.delete(nodeId); } } isOnline(nodeId: string): boolean { const conn = this.connections.get(nodeId); return !!conn && conn.socket.readyState === 1; } private send(socket: WebSocket, message: ControlMessage): void { if (socket.readyState === 1) { socket.send(JSON.stringify(message)); } } private async checkOfflineNodes(): Promise { if (!this.config) return; const threshold = new Date( Date.now() - this.config.NODE_OFFLINE_THRESHOLD_SECONDS * 1000 ); const stale = await prisma.node.findMany({ where: { status: "ONLINE", lastSeen: { lt: threshold }, revoked: false, }, select: { id: true }, }); for (const node of stale) { if (this.connections.has(node.id)) { this.disconnect(node.id); } await prisma.node.update({ where: { id: node.id }, data: { status: "OFFLINE" }, }); await prisma.mediaFile.updateMany({ where: { nodeId: node.id }, data: { available: false }, }); } } shutdown(): void { if (this.offlineCheckInterval) clearInterval(this.offlineCheckInterval); for (const conn of this.connections.values()) { conn.socket.close(1001, "Server shutting down"); } this.connections.clear(); this.syncState.clear(); } } export const nodeConnectionManager = new NodeConnectionManager();