node_server/bot/forward/manager.ts

import {createHash} from 'node:crypto';
import {io, type Socket} from 'socket.io-client';
import {BotStream} from './stream.ts';
import {BotFileService} from './files.ts';
import type {EnsureForward, ForwardAck, ForwardReady, ForwardScope} from '../../types/forward.ts';
import {error_ack, exact, failure_code, ForwardFailure, success_data, validate_command, validate_join,
    validate_rpc} from '../../utils/forward_protocol.ts';

type Entry = {
    command: EnsureForward;
    chain: Promise<unknown>;
    pending: number;
    socket?: Socket;
    detach?: () => void;
    ready: boolean;
    terminal: boolean;
    epoch: number;
    lease: number;
    retain_until: number;
    expiry?: ReturnType<typeof setTimeout>;
    stream?: BotStream;
    files?: Pick<BotFileService, 'list' | 'read' | 'dispose'>;
    unsubscribe?: () => void;
    cancel: Set<() => void>;
    used_tokens: Map<string, number>;
};

type Options = {
    socket_factory?: typeof io;
    stream_factory?: (bot_id: number) => BotStream;
    file_factory?: (bot_id: number) => Pick<BotFileService, 'list' | 'read' | 'dispose'>;
};

const scope_keys = ['forward_id', 'forward_session_id', 'node_session_id', 'node_id', 'bot_id'] as const;
const fingerprint = (token: string) => createHash('sha256').update(token).digest('hex');
const managers = new WeakMap<Socket, NodeForwardManager>();

export class NodeForwardManager {
    private readonly entries = new Map<string, Entry>();
    private readonly streams = new Map<number, {stream: BotStream; refs: number}>();
    private readonly socket_factory: typeof io;
    private readonly stream_factory: (bot_id: number) => BotStream;
    private readonly file_factory: (bot_id: number) => Pick<BotFileService, 'list' | 'read' | 'dispose'>;
    private stopped = false;

    constructor(readonly node_id: number, readonly node_session_id: string, options: Options = {}) {
        this.socket_factory = options.socket_factory ?? io;
        this.stream_factory = options.stream_factory ?? (bot_id => new BotStream(bot_id));
        this.file_factory = options.file_factory ?? (bot_id => new BotFileService(bot_id));
    }

    async ensure(raw: unknown): Promise<ForwardAck<ForwardReady>> {
        let entry: Entry | undefined;
        try {
            const command = validate_command(raw);
            this.check_command(command);
            for (const [id, previous] of this.entries) {
                if (previous.terminal && previous.pending === 0
                    && previous.retain_until + 5000 <= Date.now()) this.entries.delete(id);
            }
            entry = this.entries.get(command.forward_session_id);
            if (entry) {
                if (scope_keys.some(key => entry!.command[key] !== command[key])
                    || (['url', 'path', 'namespace'] as const).some(key => entry!.command[key] !== command[key])) {
                    throw new ForwardFailure('SESSION_STALE');
                }
                if (entry.terminal) throw new ForwardFailure('SESSION_STALE');
                if (entry.pending >= 16) throw new ForwardFailure('RATE_LIMITED');
            } else {
                if (this.entries.size >= 4096 || [...this.entries.values()].filter(item => !item.terminal).length >= 500) {
                    throw new ForwardFailure('RATE_LIMITED');
                }
                // Retain closed identities past all known credentials, separately from active capacity.
                entry = {command: {...command, token: ''}, chain: Promise.resolve(), pending: 0, ready: false, terminal: false,
                    epoch: 0, lease: 0, retain_until: command.expires_at, cancel: new Set(), used_tokens: new Map()};
                this.entries.set(command.forward_session_id, entry);
                this.arm_expiry(entry);
            }
            const current = entry;
            current.retain_until = Math.max(current.retain_until, command.expires_at);
            current.pending++;
            let running = false;
            const run = current.chain.then(() => {
                running = true;
                return this.establish(current, command);
            });
            current.chain = run.then(() => { current.pending--; }, () => { current.pending--; });
            return await this.until_deadline(run, command.deadline_at, () => {
                if (running) this.close_transport(current);
            });
        } catch (error) {
            return error_ack(failure_code(error));
        }
    }

    private check_command(command: EnsureForward) {
        if (this.stopped || command.node_session_id !== this.node_session_id) throw new ForwardFailure('SESSION_STALE');
        if (command.node_id !== this.node_id) throw new ForwardFailure('FORBIDDEN');
        if (command.deadline_at <= Date.now()) throw new ForwardFailure('TIMEOUT');
        if (command.expires_at <= Date.now()) throw new ForwardFailure('TOKEN_EXPIRED');
    }

    private scope(command: EnsureForward): ForwardScope {
        const {forward_id, forward_session_id, node_session_id, node_id, bot_id} = command;
        return {forward_id, forward_session_id, node_session_id, node_id, bot_id};
    }

    private async establish(entry: Entry, command: EnsureForward): Promise<ForwardAck<ForwardReady>> {
        this.check_command(command);
        if (entry.terminal) throw new ForwardFailure('SESSION_STALE');
        if (command.expires_at < entry.lease) throw new ForwardFailure('SESSION_STALE');
        const digest = fingerprint(command.token);
        for (const [token, expiry] of entry.used_tokens) if (expiry <= Date.now()) entry.used_tokens.delete(token);
        if (!entry.used_tokens.has(digest) && entry.used_tokens.size >= 256) throw new ForwardFailure('RATE_LIMITED');
        let socket = entry.socket;
        try {
            if (socket?.connected && entry.ready) {
                entry.used_tokens.set(digest, command.expires_at);
                const joined = validate_join(await this.ack(entry, socket, 'reauth', {token: command.token}, command), command, false);
                this.check_command(command);
                if (!joined.node_ready) {
                    entry.ready = false;
                    await this.ready(entry, socket, command);
                }
            } else {
                if (entry.used_tokens.has(digest)) throw new ForwardFailure('TOKEN_EXPIRED');
                this.close_transport(entry);
                // No autonomous reconnect: a join window can expire before the lease does.
                socket = this.socket_factory(`${command.url.replace(/\/$/, '')}${command.namespace === '/' ? '' : command.namespace}`, {
                    path: command.path, autoConnect: false, reconnection: false, forceNew: true,
                    timeout: Math.min(10_000, this.remaining(command)), auth: {token: command.token},
                });
                entry.socket = socket;
                this.install_transport(entry, socket);
                await this.connect(entry, socket, command);
                entry.used_tokens.set(digest, command.expires_at);
                validate_join(await this.ack(entry, socket, 'join_room', {token: command.token}, command), command, true);
                this.check_command(command);
                if (!entry.stream) {
                    let shared = this.streams.get(command.bot_id);
                    if (!shared) {
                        shared = {stream: this.stream_factory(command.bot_id), refs: 0};
                        this.streams.set(command.bot_id, shared);
                    }
                    shared.refs++;
                    entry.stream = shared.stream;
                }
                this.observe_updates(entry, socket);
                await this.ready(entry, socket, command);
            }
            this.check_command(command);
            if (entry.socket !== socket || !socket.connected || entry.terminal || !entry.ready) throw new ForwardFailure('SESSION_STALE');
            entry.command = {...command, token: ''};
            entry.lease = command.expires_at;
            entry.used_tokens.set(digest, command.expires_at);
            this.arm_expiry(entry);
            entry.stream?.resync();
            this.check_command(command);
            return {status: 'success', data: {...this.scope(command), node_ready: true, lease_expires_at: command.expires_at}};
        } catch (error) {
            this.close_transport(entry);
            if (!entry.lease) this.release_stream(entry);
            throw error;
        }
    }

    private remaining(command: EnsureForward) {
        const remaining = Math.min(command.deadline_at, command.expires_at) - Date.now();
        if (remaining <= 0) throw new ForwardFailure('TIMEOUT');
        return Math.min(remaining, 2_147_483_647);
    }

    private until_deadline<T>(task: Promise<T>, deadline: number, timeout: () => void): Promise<T> {
        return new Promise((resolve, reject) => {
            const timer = setTimeout(() => {
                timeout();
                reject(new ForwardFailure('TIMEOUT'));
            }, Math.max(0, Math.min(deadline - Date.now(), 2_147_483_647)));
            task.then(resolve, reject).finally(() => clearTimeout(timer));
        });
    }

    private connect(entry: Entry, socket: Socket, command: EnsureForward): Promise<void> {
        return new Promise((resolve, reject) => {
            const done = (error?: ForwardFailure) => {
                clearTimeout(timer);
                socket.off('connect', connected);
                socket.off('connect_error', failed);
                entry.cancel.delete(cancel);
                if (error) reject(error); else resolve();
            };
            const connected = () => done();
            const failed = () => done(new ForwardFailure('RELAY_UNAVAILABLE'));
            const cancel = () => done(new ForwardFailure('CANCELLED'));
            const timer = setTimeout(() => done(new ForwardFailure('TIMEOUT')), this.remaining(command));
            entry.cancel.add(cancel);
            socket.once('connect', connected);
            socket.once('connect_error', failed);
            socket.connect();
        });
    }

    private async ack(entry: Entry, socket: Socket, event: string, payload: unknown, command: EnsureForward): Promise<unknown> {
        if (entry.terminal || entry.socket !== socket || !socket.connected || this.stopped) throw new ForwardFailure('SESSION_STALE');
        const epoch = entry.epoch;
        const result = await new Promise<unknown>((resolve, reject) => {
            const cancel = () => { entry.cancel.delete(cancel); reject(new ForwardFailure('CANCELLED')); };
            entry.cancel.add(cancel);
            socket.timeout(this.remaining(command)).emit(event, payload, (error: unknown, response: unknown) => {
                entry.cancel.delete(cancel);
                if (error) reject(new ForwardFailure('TIMEOUT')); else resolve(response);
            });
        });
        if (entry.terminal || entry.epoch !== epoch || entry.socket !== socket || !socket.connected || this.stopped) {
            throw new ForwardFailure('SESSION_STALE');
        }
        return result;
    }

    private async ready(entry: Entry, socket: Socket, command: EnsureForward) {
        const data = success_data(await this.ack(entry, socket, 'forward_node_ready', {
            forward_session_id: command.forward_session_id,
        }, command));
        if (!exact(data, ['node_ready']) || data.node_ready !== true) throw new ForwardFailure('RELAY_UNAVAILABLE');
        entry.ready = true;
    }

    private install_transport(entry: Entry, socket: Socket) {
        const seen = new Map<string, number>();
        const rpc = async (raw: unknown, callback: unknown) => {
            if (typeof callback !== 'function') return;
            let response: ForwardAck;
            try {
                // Relay may dispatch immediately after accepting ready, before its ACK reaches Node.
                if (entry.socket !== socket || !socket.connected || entry.terminal || this.stopped || !entry.stream) {
                    throw new ForwardFailure('SESSION_STALE');
                }
                const request = validate_rpc(raw);
                if (request.forward_session_id !== entry.command.forward_session_id) throw new ForwardFailure('SESSION_STALE');
                if (request.deadline_at <= Date.now()) throw new ForwardFailure('TIMEOUT');
                if (entry.lease && entry.lease <= Date.now()) throw new ForwardFailure('TOKEN_EXPIRED');
                for (const [id, deadline] of seen) if (deadline <= Date.now()) seen.delete(id);
                if (seen.has(request.correlation_id)) throw new ForwardFailure('CANCELLED');
                if (seen.size >= 512) throw new ForwardFailure('RATE_LIMITED');
                seen.set(request.correlation_id, request.deadline_at);
                let data: unknown;
                switch (request.method) {
                    case 'bot.snapshot': data = entry.stream.snapshot(); break;
                    case 'get_context_info': data = entry.stream.get_context_info(request.payload, request.deadline_at); break;
                    case 'bot.logs': data = entry.stream.logs(request.payload as {limit: number; stream_id?: string; after_seq?: number}); break;
                    case 'get_log_history':
                    case 'get_data_file_list':
                    case 'get_file_data':
                    case 'execute_function':
                    case 'world.snapshot':
                    case 'world.chunk':
                    case 'world.unwatch':
                    case 'workflow.start':
                    case 'workflow.stop':
                    case 'workflow.snapshot':
                    case 'workflow.catalog':
                    case 'workflow.runtime':
                    case 'workflow.control':
                    case 'workflow.command':
                    case 'tab_complete': {
                        const epoch = entry.epoch;
                        const check_transport = () => {
                            if (entry.socket !== socket || !socket.connected || entry.terminal || this.stopped || entry.epoch !== epoch)
                                throw new ForwardFailure('SESSION_STALE');
                            if (request.deadline_at <= Date.now()) throw new ForwardFailure('TIMEOUT');
                            if (entry.lease && entry.lease <= Date.now()) throw new ForwardFailure('TOKEN_EXPIRED');
                        };
                        if (request.method === 'get_data_file_list' || request.method === 'get_file_data') {
                            // File IDs belong to this transport/session, never the shared BotStream.
                            const files = entry.files ??= this.file_factory(entry.command.bot_id);
                            const deadline = Math.min(request.deadline_at, entry.lease || entry.command.expires_at);
                            data = request.method === 'get_data_file_list' ? await files.list(deadline, entry.lease || entry.command.expires_at)
                                : await files.read(request.payload as {id: string; offset?: number; length?: number}, deadline);
                            if (request.method === 'get_data_file_list') {
                                const listing = data as {expires_at: number};
                                listing.expires_at = Math.min(listing.expires_at, entry.lease || entry.command.expires_at);
                            }
                        } else if (request.method === 'world.snapshot') {
                            data = await this.until_deadline(entry.stream.world_snapshot(request.payload, request.deadline_at, check_transport),
                                request.deadline_at, () => {});
                        } else if (request.method === 'world.chunk') {
                            data = await this.until_deadline(entry.stream.world_chunk(request.payload, request.deadline_at, check_transport),
                                request.deadline_at, () => {});
                        } else if (request.method === 'world.unwatch') {
                            data = entry.stream.world_unwatch(request.payload, request.deadline_at, check_transport);
                        } else if (request.method === 'workflow.start') {
                            data = entry.stream.workflow_start(request.payload, request.deadline_at, check_transport);
                        } else if (request.method === 'workflow.stop') {
                            data = await this.until_deadline(entry.stream.workflow_stop(request.payload, request.deadline_at, check_transport),
                                request.deadline_at, () => {});
                        } else if (request.method === 'workflow.snapshot') {
                            data = entry.stream.workflow_snapshot(request.payload, request.deadline_at, check_transport);
                        } else if (request.method === 'workflow.catalog') {
                            data = entry.stream.workflow_catalog(request.payload, request.deadline_at, check_transport);
                        } else if (request.method === 'workflow.runtime') {
                            data = entry.stream.workflow_runtime(request.payload, request.deadline_at, check_transport);
                        } else if (request.method === 'workflow.control') {
                            data = entry.stream.workflow_control(request.payload, request.deadline_at, check_transport);
                        } else if (request.method === 'workflow.command') {
                            data = entry.stream.workflow_command(request.payload, request.deadline_at, check_transport);
                        } else data = request.method === 'get_log_history' ? await entry.stream.get_log_history()
                            : request.method === 'execute_function' ? await entry.stream.execute_function(
                                request.payload as {function: string; name: string; args: unknown[]}, request.deadline_at, check_transport)
                            : await entry.stream.tab_complete(request.payload, request.deadline_at);
                        if (entry.socket !== socket || !socket.connected || entry.terminal || this.stopped || entry.epoch !== epoch)
                            throw new ForwardFailure('SESSION_STALE');
                        if (request.deadline_at <= Date.now()) throw new ForwardFailure('TIMEOUT');
                        if (entry.lease && entry.lease <= Date.now()) throw new ForwardFailure('TOKEN_EXPIRED');
                        break;
                    }
                    case 'bot.chat': data = entry.stream.chat(request.payload as {message: string}, request.deadline_at); break;
                }
                response = {status: 'success', data};
            } catch (error) {
                response = error_ack(failure_code(error));
            }
            try { callback(response); } catch { /* The request ACK must never escape an event callback. */ }
        };
        const closed = (raw: unknown) => {
            if (entry.socket !== socket || !exact(raw, ['reason', 'code'], ['forward_session_id'])
                || (raw.forward_session_id !== undefined && raw.forward_session_id !== entry.command.forward_session_id)) return;
            this.retire(entry);
        };
        const disconnected = () => {
            if (entry.socket === socket) this.close_transport(entry);
        };
        const status = (raw: unknown) => {
            if (entry.socket === socket && exact(raw, ['forward_session_id', 'node_ready', 'reason'])
                && raw.forward_session_id === entry.command.forward_session_id && raw.reason === 'resync_required') {
                entry.stream?.resync();
            }
        };
        socket.on('relay_request', rpc);
        socket.on('forward_session_closed', closed);
        socket.on('relay_status', status);
        socket.on('disconnect', disconnected);
        entry.detach = () => {
            socket.off('relay_request', rpc);
            socket.off('forward_session_closed', closed);
            socket.off('relay_status', status);
            socket.off('disconnect', disconnected);
        };
    }

    private observe_updates(entry: Entry, socket: Socket) {
        entry.unsubscribe?.();
        let window = Date.now();
        let count = 0;
        entry.unsubscribe = entry.stream!.subscribe(event => {
            if ((event.name !== 'log' && event.name !== 'info_update' && event.name !== 'status_update') || entry.socket !== socket || !socket.connected
                || !entry.ready || entry.terminal || this.stopped) return;
            if (event.name === 'status_update') {
                // Lifecycle state must survive a busy log stream; reconnect rather
                // than silently retaining an online snapshot when the queue is full.
                if ((socket.io.engine?.writeBuffer.length ?? 0) >= 64) this.close_transport(entry);
                else socket.emit('bot_event', {...event, forward_session_id: entry.command.forward_session_id});
                return;
            }
            if (Date.now() - window >= 1000) { window = Date.now(); count = 0; }
            if (++count > 150 || (socket.io.engine?.writeBuffer.length ?? 0) >= 64) return;
            if (event.name === 'log') {
                socket.volatile.emit('bot_log', {...event, forward_session_id: entry.command.forward_session_id});
            } else {
                const {name, payload, ...metadata} = event;
                const update = {...metadata,
                    forward_session_id: entry.command.forward_session_id, data: payload};
                // Plugin lifecycle transitions can finish within the same transport
                // write. Volatile would drop the final state while that write drains,
                // leaving the UI initializing until another plugin change occurs.
                // Keep the shared rate/buffer bounds above for these queued packets.
                if (payload.name === 'plugin_states' || payload.name === 'workflow_states') socket.emit('info_update', update);
                else socket.volatile.emit('info_update', update);
            }
        });
    }

    private close_transport(entry: Entry) {
        entry.epoch++;
        entry.ready = false;
        entry.unsubscribe?.();
        entry.unsubscribe = undefined;
        entry.files?.dispose();
        entry.files = undefined;
        const socket = entry.socket;
        entry.socket = undefined;
        entry.detach?.();
        entry.detach = undefined;
        for (const cancel of [...entry.cancel]) cancel();
        entry.cancel.clear();
        socket?.disconnect();
        if (socket) socket.auth = {};
    }

    private arm_expiry(entry: Entry) {
        clearTimeout(entry.expiry);
        const expires_at = Math.max(entry.lease, entry.command.expires_at);
        entry.expiry = setTimeout(() => {
            if (expires_at <= Date.now()) this.retire(entry);
            else this.arm_expiry(entry);
        }, Math.min(Math.max(0, expires_at - Date.now()), 2_147_483_647));
        entry.expiry.unref();
    }

    private release_stream(entry: Entry) {
        if (!entry.stream) return;
        const shared = this.streams.get(entry.command.bot_id);
        if (shared && --shared.refs === 0) {
            shared.stream.dispose();
            this.streams.delete(entry.command.bot_id);
        }
        entry.stream = undefined;
    }

    private retire(entry: Entry) {
        entry.terminal = true;
        clearTimeout(entry.expiry);
        this.close_transport(entry);
        this.release_stream(entry);
        entry.used_tokens.clear();
    }

    stop() {
        if (this.stopped) return;
        this.stopped = true;
        for (const entry of this.entries.values()) this.retire(entry);
        this.entries.clear();
    }
}

export function bind_node_session(socket: Socket, node_id: number, node_session_id: string): void {
    clear_node_session(socket);
    managers.set(socket, new NodeForwardManager(node_id, node_session_id));
}

export function clear_node_session(socket: Socket): void {
    managers.get(socket)?.stop();
    managers.delete(socket);
}

export function ensure_forward(socket: Socket, command: unknown): Promise<ForwardAck<ForwardReady>> {
    if (!socket.connected) return Promise.resolve(error_ack('NODE_OFFLINE'));
    return managers.get(socket)?.ensure(command) ?? Promise.resolve(error_ack('UNAUTHORIZED'));
}